From e533e30851cd8adb1a774e0363d932314f3fcc95 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 14:06:28 +0800 Subject: [PATCH 01/16] rwkv: fuse wkv7 graph operations --- ggml/include/ggml.h | 3 +- ggml/src/ggml-cpu/ops.cpp | 88 ++++++--- ggml/src/ggml-cuda/binbcast.cu | 168 ++++++++++++++++++ ggml/src/ggml-cuda/binbcast.cuh | 14 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 162 +++++++++++++++++ ggml/src/ggml-cuda/wkv.cu | 58 +++--- ggml/src/ggml-metal/ggml-metal-device.cpp | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 3 +- ggml/src/ggml-metal/ggml-metal.metal | 34 ++-- ggml/src/ggml-sycl/wkv.cpp | 53 +++--- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 12 +- ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp | 35 ++-- ggml/src/ggml.c | 20 ++- src/models/rwkv7-base.cpp | 11 +- tests/test-backend-ops.cpp | 7 +- 15 files changed, 553 insertions(+), 117 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d6807b6dd47a..f3839e1a5688 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2525,8 +2525,9 @@ extern "C" { struct ggml_tensor * w, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * kk, struct ggml_tensor * a, - struct ggml_tensor * b, + struct ggml_tensor * r_k, struct ggml_tensor * state); /* Solves a specific equation of the form Ax=B, where A is a triangular matrix diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 6724686b8ae2..afbc98ffc1d6 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11,6 +11,7 @@ #include #include #include +#include // ggml_compute_forward_dup @@ -10825,17 +10826,22 @@ void ggml_compute_forward_gated_delta_net( // ggml_compute_forward_rwkv_wkv7 +static inline float ggml_rwkv_wkv7_decay(float w) { + constexpr float w_scale = -0.6065306597126334f; + return expf(w_scale / (1.0f + expf(-w))); +} + static void ggml_compute_forward_rwkv_wkv7_f32( const ggml_compute_params * params, ggml_tensor * dst) { const int64_t T = dst->src[1]->ne[2]; const int64_t C = dst->ne[0]; const int64_t HEADS = dst->src[1]->ne[1]; - const int64_t n_seqs = dst->src[6]->ne[1]; + const int64_t n_seqs = dst->src[7]->ne[1]; const int64_t head_size = C / HEADS; float * dst_data = (float *) dst->data; - float * state = ((float *) dst->data) + C * T; + float * state = ((float *) dst->data) + 2 * C * T; const int ith = params->ith; const int nth = params->nth; @@ -10848,14 +10854,16 @@ static void ggml_compute_forward_rwkv_wkv7_f32( float * w = (float *) dst->src[1]->data; float * k = (float *) dst->src[2]->data; float * v = (float *) dst->src[3]->data; - float * a = (float *) dst->src[4]->data; - float * b = (float *) dst->src[5]->data; + float * kk = (float *) dst->src[4]->data; + float * a = (float *) dst->src[5]->data; + float * r_k = (float *) dst->src[6]->data; int64_t t_stride = HEADS * head_size; // Same to C int64_t h_stride = C / HEADS; GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS int64_t h_stride_2d = head_size * head_size; + std::vector w_decay(head_size); #if defined(GGML_SIMD) #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v_intrinsic) @@ -10864,13 +10872,20 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; + float rk = 0; + for (int64_t j = 0; j < head_size; j++) { + const int64_t t_h_j_offset = t_h_offset + j; + rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; + w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + } + for (int64_t i = 0; i < head_size; i++) { int64_t t_h_i_offset = t_h_offset + i; int64_t h_2d_i_offset = h_2d_offset + i * h_stride; @@ -10879,23 +10894,25 @@ static void ggml_compute_forward_rwkv_wkv7_f32( float sa = 0, result = 0; for (int64_t j = 0; j < head_size; j++) { - sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; + sa += kk[t_h_offset + j] * state_prev[h_2d_i_offset + j]; } + sa = -sa; for (int64_t j = 0; j < head_size; j++) { int64_t t_h_j_offset = t_h_offset + j; int64_t h_2d_i_j_offset = h_2d_i_offset + j; float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; + float w_val = w_decay[j]; float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; + float b_val = kk[t_h_j_offset] * a[t_h_j_offset]; float kv_val = v_val * k_val; float prev_state_val = state_prev[h_2d_i_j_offset]; state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; result += state_cur[h_2d_i_j_offset] * r_val; } dst_data[t_h_i_offset] = result; + dst_data[C * T + t_h_i_offset] = v_val * rk; } } } @@ -10904,13 +10921,20 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; + float rk = 0; + for (int64_t j = 0; j < head_size; j++) { + const int64_t t_h_j_offset = t_h_offset + j; + rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; + w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + } + for (int64_t ii = 0; ii < head_size; ii++) { int64_t t_h_i_offset = t_h_offset + ii; int64_t h_2d_i_offset = h_2d_offset + ii * h_stride; @@ -10923,13 +10947,14 @@ static void ggml_compute_forward_rwkv_wkv7_f32( GGML_F32_VEC ax[GGML_F32_ARR]; GGML_F32_VEC ay[GGML_F32_ARR]; for (int64_t j = 0; j < head_size; j += GGML_F32_STEP) { - for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { - ax[kk] = GGML_F32_VEC_LOAD(&a[t_h_offset + j + kk * GGML_F32_EPR]); - ay[kk] = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_offset + j + kk * GGML_F32_EPR]); - sum[kk] = GGML_F32_VEC_FMA(sum[kk], ax[kk], ay[kk]); + for (int64_t iv = 0; iv < GGML_F32_ARR; iv++) { + ax[iv] = GGML_F32_VEC_LOAD(&kk[t_h_offset + j + iv * GGML_F32_EPR]); + ay[iv] = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_offset + j + iv * GGML_F32_EPR]); + sum[iv] = GGML_F32_VEC_FMA(sum[iv], ax[iv], ay[iv]); } } GGML_F32_VEC_REDUCE(sa, sum); + sa = -sa; } GGML_F32_VEC sa_vec = GGML_F32_VEC_SET1(sa); @@ -10937,14 +10962,16 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t j = 0; GGML_F32_VEC result_vec[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; for (; j < head_size; j += GGML_F32_STEP) { - for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { - int64_t t_h_j_offset = t_h_offset + j + kk * GGML_F32_EPR; - int64_t h_2d_i_j_offset = h_2d_i_offset + j + kk * GGML_F32_EPR; + for (int64_t iv = 0; iv < GGML_F32_ARR; iv++) { + int64_t t_h_j_offset = t_h_offset + j + iv * GGML_F32_EPR; + int64_t h_2d_i_j_offset = h_2d_i_offset + j + iv * GGML_F32_EPR; GGML_F32_VEC r_vec = GGML_F32_VEC_LOAD(&r[t_h_j_offset]); - GGML_F32_VEC w_vec = GGML_F32_VEC_LOAD(&w[t_h_j_offset]); + GGML_F32_VEC w_vec = GGML_F32_VEC_LOAD(&w_decay[j + iv * GGML_F32_EPR]); GGML_F32_VEC k_vec = GGML_F32_VEC_LOAD(&k[t_h_j_offset]); - GGML_F32_VEC b_vec = GGML_F32_VEC_LOAD(&b[t_h_j_offset]); + GGML_F32_VEC kk_vec = GGML_F32_VEC_LOAD(&kk[t_h_j_offset]); + GGML_F32_VEC a_vec = GGML_F32_VEC_LOAD(&a[t_h_j_offset]); + GGML_F32_VEC b_vec = GGML_F32_VEC_MUL(kk_vec, a_vec); k_vec = GGML_F32_VEC_MUL(v_vec, k_vec); @@ -10954,10 +10981,11 @@ static void ggml_compute_forward_rwkv_wkv7_f32( state_vec = GGML_F32_VEC_FMA(state_vec, sa_vec, b_vec); GGML_F32_VEC_STORE(&state_cur[h_2d_i_j_offset], state_vec); - result_vec[kk] = GGML_F32_VEC_FMA(result_vec[kk], state_vec, r_vec); + result_vec[iv] = GGML_F32_VEC_FMA(result_vec[iv], state_vec, r_vec); } } GGML_F32_VEC_REDUCE(dst_data[t_h_i_offset], result_vec); + dst_data[C * T + t_h_i_offset] = v[t_h_i_offset] * rk; // There shouldn't be left-overs though. for (; j < head_size; j++) { @@ -10965,14 +10993,15 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t h_2d_i_j_offset = h_2d_i_offset + j; float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; + float w_val = w_decay[j]; float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; + float b_val = kk[t_h_j_offset] * a[t_h_j_offset]; float kv_val = v[t_h_i_offset] * k_val; float prev_state_val = state_prev[h_2d_i_j_offset]; state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; dst_data[t_h_i_offset] += state_cur[h_2d_i_j_offset] * r_val; + dst_data[C * T + t_h_i_offset] = v[t_h_i_offset] * rk; } } } @@ -10983,13 +11012,20 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; + float rk = 0; + for (int64_t j = 0; j < head_size; j++) { + const int64_t t_h_j_offset = t_h_offset + j; + rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; + w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + } + for (int64_t i = 0; i < head_size; i++) { int64_t t_h_i_offset = t_h_offset + i; int64_t h_2d_i_offset = h_2d_offset + i * h_stride; @@ -10998,23 +11034,25 @@ static void ggml_compute_forward_rwkv_wkv7_f32( float sa = 0, result = 0; for (int64_t j = 0; j < head_size; j++) { - sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; + sa += kk[t_h_offset + j] * state_prev[h_2d_i_offset + j]; } + sa = -sa; for (int64_t j = 0; j < head_size; j++) { int64_t t_h_j_offset = t_h_offset + j; int64_t h_2d_i_j_offset = h_2d_i_offset + j; float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; + float w_val = w_decay[j]; float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; + float b_val = kk[t_h_j_offset] * a[t_h_j_offset]; float kv_val = v_val * k_val; float prev_state_val = state_prev[h_2d_i_j_offset]; state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; result += state_cur[h_2d_i_j_offset] * r_val; } dst_data[t_h_i_offset] = result; + dst_data[C * T + t_h_i_offset] = v_val * rk; } } } diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 2e38077bf67f..6e11dc6e71af 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -26,6 +26,86 @@ static __device__ __forceinline__ float op_div(const float a, const float b) { return a / b; } +static __device__ __forceinline__ size_t ggml_cuda_index_4d( + const int64_t i0, const int64_t i1, const int64_t i2, const int64_t i3, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s0, const int64_t s1, const int64_t s2, const int64_t s3) { + const int64_t j0 = i0 % ne0; + const int64_t j1 = i1 % ne1; + const int64_t j2 = i2 % ne2; + const int64_t j3 = i3 % ne3; + + return size_t(j0)*s0 + size_t(j1)*s1 + size_t(j2)*s2 + size_t(j3)*s3; +} + +static __global__ void k_lerp_fused_f32( + const float * __restrict__ x_prev, + const float * __restrict__ cur, + const float * __restrict__ weight, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t xp_ne0, const int64_t xp_ne1, const int64_t xp_ne2, const int64_t xp_ne3, + const int64_t xp_s0, const int64_t xp_s1, const int64_t xp_s2, const int64_t xp_s3, + const int64_t c_ne0, const int64_t c_ne1, const int64_t c_ne2, const int64_t c_ne3, + const int64_t c_s0, const int64_t c_s1, const int64_t c_s2, const int64_t c_s3, + const int64_t w_ne0, const int64_t w_ne1, const int64_t w_ne2, const int64_t w_ne3, + const int64_t w_s0, const int64_t w_s1, const int64_t w_s2, const int64_t w_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t ixp = ggml_cuda_index_4d(i0, i1, i2, i3, xp_ne0, xp_ne1, xp_ne2, xp_ne3, xp_s0, xp_s1, xp_s2, xp_s3); + const size_t ic = ggml_cuda_index_4d(i0, i1, i2, i3, c_ne0, c_ne1, c_ne2, c_ne3, c_s0, c_s1, c_s2, c_s3); + const size_t iw = ggml_cuda_index_4d(i0, i1, i2, i3, w_ne0, w_ne1, w_ne2, w_ne3, w_s0, w_s1, w_s2, w_s3); + + const float c = cur[ic]; + dst[idx] = c + (x_prev[ixp] - c) * weight[iw]; +} + +static __global__ void k_mul_sub_add_fused_f32( + const float * __restrict__ base, + const float * __restrict__ scale, + const float * __restrict__ value, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t b_ne0, const int64_t b_ne1, const int64_t b_ne2, const int64_t b_ne3, + const int64_t b_s0, const int64_t b_s1, const int64_t b_s2, const int64_t b_s3, + const int64_t s_ne0, const int64_t s_ne1, const int64_t s_ne2, const int64_t s_ne3, + const int64_t s_s0, const int64_t s_s1, const int64_t s_s2, const int64_t s_s3, + const int64_t v_ne0, const int64_t v_ne1, const int64_t v_ne2, const int64_t v_ne3, + const int64_t v_s0, const int64_t v_s1, const int64_t v_s2, const int64_t v_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t ib = ggml_cuda_index_4d(i0, i1, i2, i3, b_ne0, b_ne1, b_ne2, b_ne3, b_s0, b_s1, b_s2, b_s3); + const size_t is = ggml_cuda_index_4d(i0, i1, i2, i3, s_ne0, s_ne1, s_ne2, s_ne3, s_s0, s_s1, s_s2, s_s3); + const size_t iv = ggml_cuda_index_4d(i0, i1, i2, i3, v_ne0, v_ne1, v_ne2, v_ne3, v_s0, v_s1, v_s2, v_s3); + + const float v = value[iv]; + dst[idx] = base[ib] + (scale[is] * v - v); +} + template type == GGML_TYPE_F32); + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(weight->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(x_prev, dst)); + GGML_ASSERT(ggml_can_repeat(cur, dst)); + GGML_ASSERT(ggml_can_repeat(weight, dst)); + GGML_ASSERT(ggml_is_contiguous(x_prev)); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(weight)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_lerp_fused_f32, launch_params, + (const float *) x_prev->data, + (const float *) cur->data, + (const float *) weight->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + x_prev->ne[0], x_prev->ne[1], x_prev->ne[2], x_prev->ne[3], + stride(x_prev, 0), stride(x_prev, 1), stride(x_prev, 2), stride(x_prev, 3), + cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3], + stride(cur, 0), stride(cur, 1), stride(cur, 2), stride(cur, 3), + weight->ne[0], weight->ne[1], weight->ne[2], weight->ne[3], + stride(weight, 0), stride(weight, 1), stride(weight, 2), stride(weight, 3)); +} + +void ggml_cuda_op_mul_sub_add_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * base, + const ggml_tensor * scale, + const ggml_tensor * value, + ggml_tensor * dst) { + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(value->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(base, dst)); + GGML_ASSERT(ggml_can_repeat(scale, dst)); + GGML_ASSERT(ggml_can_repeat(value, dst)); + GGML_ASSERT(ggml_is_contiguous(base)); + GGML_ASSERT(ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(value)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_mul_sub_add_fused_f32, launch_params, + (const float *) base->data, + (const float *) scale->data, + (const float *) value->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + base->ne[0], base->ne[1], base->ne[2], base->ne[3], + stride(base, 0), stride(base, 1), stride(base, 2), stride(base, 3), + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], + stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3), + value->ne[0], value->ne[1], value->ne[2], value->ne[3], + stride(value, 0), stride(value, 1), stride(value, 2), stride(value, 3)); +} + void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 12624785b444..b570401cc286 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -10,3 +10,17 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); + +void ggml_cuda_op_lerp_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * x_prev, + const ggml_tensor * cur, + const ggml_tensor * weight, + ggml_tensor * dst); + +void ggml_cuda_op_mul_sub_add_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * base, + const ggml_tensor * scale, + const ggml_tensor * value, + ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index cca70592f807..dead93f74765 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3549,6 +3549,120 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod return true; } +static bool ggml_cuda_should_fuse_lerp( + const ggml_tensor * sub, + const ggml_tensor * repeat, + const ggml_tensor * mul, + const ggml_tensor * add, + const ggml_tensor ** x_prev, + const ggml_tensor ** cur, + const ggml_tensor ** weight) { + if (sub->op != GGML_OP_SUB || mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD) { + return false; + } + + const ggml_tensor * sx = sub; + if (repeat) { + if (repeat->op != GGML_OP_REPEAT || repeat->src[0] != sub) { + return false; + } + sx = repeat; + } + + const ggml_tensor * w = nullptr; + if (mul->src[0] == sx) { + w = mul->src[1]; + } else if (mul->src[1] == sx) { + w = mul->src[0]; + } else { + return false; + } + + const ggml_tensor * add_cur = nullptr; + if (add->src[0] == mul) { + add_cur = add->src[1]; + } else if (add->src[1] == mul) { + add_cur = add->src[0]; + } else { + return false; + } + + if (add_cur != sub->src[1]) { + return false; + } + + const ggml_tensor * xp = sub->src[0]; + const ggml_tensor * c = sub->src[1]; + + if (xp->type != GGML_TYPE_F32 || c->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_can_repeat(xp, add) || !ggml_can_repeat(c, add) || !ggml_can_repeat(w, add)) { + return false; + } + + if (!ggml_is_contiguous(xp) || !ggml_is_contiguous(c) || !ggml_is_contiguous(w) || !ggml_is_contiguous(add)) { + return false; + } + + *x_prev = xp; + *cur = c; + *weight = w; + return true; +} + +static bool ggml_cuda_should_fuse_mul_sub_add( + const ggml_tensor * mul, + const ggml_tensor * sub, + const ggml_tensor * add, + const ggml_tensor ** base, + const ggml_tensor ** scale, + const ggml_tensor ** value) { + if (mul->op != GGML_OP_MUL || sub->op != GGML_OP_SUB || add->op != GGML_OP_ADD) { + return false; + } + if (sub->src[0] != mul) { + return false; + } + + const ggml_tensor * v = sub->src[1]; + const ggml_tensor * s = nullptr; + if (mul->src[0] == v) { + s = mul->src[1]; + } else if (mul->src[1] == v) { + s = mul->src[0]; + } else { + return false; + } + + const ggml_tensor * b = nullptr; + if (add->src[0] == sub) { + b = add->src[1]; + } else if (add->src[1] == sub) { + b = add->src[0]; + } else { + return false; + } + + if (b->type != GGML_TYPE_F32 || s->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_can_repeat(b, add) || !ggml_can_repeat(s, add) || !ggml_can_repeat(v, add)) { + return false; + } + + if (!ggml_is_contiguous(b) || !ggml_is_contiguous(s) || !ggml_is_contiguous(v) || !ggml_is_contiguous(add)) { + return false; + } + + *base = b; + *scale = s; + *value = v; + return true; +} + // returns whether the write (out) nodes overwrite the read nodes in operation static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, @@ -3923,6 +4037,54 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } + // RWKV lerp: cur + (x_prev - cur) * weight. Time-mix has an extra REPEAT + // between SUB and MUL; channel-mix uses SUB directly. + if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_REPEAT, GGML_OP_MUL, GGML_OP_ADD }, { i + 3 })) { + const ggml_tensor * x_prev = nullptr; + const ggml_tensor * cur = nullptr; + const ggml_tensor * weight = nullptr; + + if (ggml_cuda_should_fuse_lerp(cgraph->nodes[i], cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], + &x_prev, &cur, &weight)) { + const int out_nodes[] = { i + 3 }; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 4, out_nodes, 1)) { + ggml_cuda_op_lerp_fused(*cuda_ctx, x_prev, cur, weight, cgraph->nodes[i + 3]); + return 3; + } + } + } + + if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_MUL, GGML_OP_ADD }, { i + 2 })) { + const ggml_tensor * x_prev = nullptr; + const ggml_tensor * cur = nullptr; + const ggml_tensor * weight = nullptr; + + if (ggml_cuda_should_fuse_lerp(cgraph->nodes[i], nullptr, cgraph->nodes[i + 1], cgraph->nodes[i + 2], + &x_prev, &cur, &weight)) { + const int out_nodes[] = { i + 2 }; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { + ggml_cuda_op_lerp_fused(*cuda_ctx, x_prev, cur, weight, cgraph->nodes[i + 2]); + return 2; + } + } + } + + // RWKV key update: k + (a * ka - ka). + if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_SUB, GGML_OP_ADD }, { i + 2 })) { + const ggml_tensor * base = nullptr; + const ggml_tensor * scale = nullptr; + const ggml_tensor * value = nullptr; + + if (ggml_cuda_should_fuse_mul_sub_add(cgraph->nodes[i], cgraph->nodes[i + 1], cgraph->nodes[i + 2], + &base, &scale, &value)) { + const int out_nodes[] = { i + 2 }; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { + ggml_cuda_op_mul_sub_add_fused(*cuda_ctx, base, scale, value, cgraph->nodes[i + 2]); + return 2; + } + } + } + // Snake activation: y = x + sin(a*x)^2 * inv_b // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add if (ggml_can_fuse_subgraph(cgraph, i, diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index d2fced705e09..3c104bfe2c23 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -66,7 +66,7 @@ static __global__ void rwkv_wkv_f32(const int B, const int T, const int C, const } template -static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * a, const float * b, const float * s, float * dst) { +static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * r_k, const float * s, float * dst) { const int tid = threadIdx.x; const int bid = blockIdx.x; @@ -77,7 +77,7 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons const int n_seq_tokens = T / B; float state[head_size]; - __shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size]; + __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size], _r_k[head_size]; #ifndef GGML_USE_MUSA #pragma unroll @@ -85,27 +85,38 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons for (int i = 0; i < head_size; i++) { state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i]; } + _r_k[tid] = r_k[head_i * head_size + tid]; for (int t = batch_i * n_seq_tokens * C + head_i * head_size + tid; t < (batch_i + 1) * n_seq_tokens * C + head_i * head_size + tid; t += C) { __syncthreads(); + constexpr float w_scale = -0.6065306597126334f; _r[tid] = r[t]; - _w[tid] = w[t]; + _w[tid] = __expf(w_scale / (1.0f + __expf(-w[t]))); _k[tid] = k[t]; + _kk[tid] = kk[t]; _a[tid] = a[t]; - _b[tid] = b[t]; __syncthreads(); float sa = 0; + float rk = 0; #pragma unroll for (int j = 0; j < head_size; j += 4) { - const float4& a = (float4&)(_a[j]); + const float4& r = (float4&)(_r[j]); + const float4& k = (float4&)(_k[j]); + const float4& kk = (float4&)(_kk[j]); + const float4& r_k = (float4&)(_r_k[j]); const float4& s = (float4&)(state[j]); - sa += a.x * s.x; - sa += a.y * s.y; - sa += a.z * s.z; - sa += a.w * s.w; + sa += kk.x * s.x; + sa += kk.y * s.y; + sa += kk.z * s.z; + sa += kk.w * s.w; + rk += r.x * k.x * r_k.x; + rk += r.y * k.y * r_k.y; + rk += r.z * k.z * r_k.z; + rk += r.w * k.w * r_k.w; } + sa = -sa; const float _v = v[t]; float y = 0; @@ -113,7 +124,8 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons const float4& r = (float4&)(_r[j]); const float4& w = (float4&)(_w[j]); const float4& k = (float4&)(_k[j]); - const float4& b = (float4&)(_b[j]); + const float4& kk = (float4&)(_kk[j]); + const float4& a = (float4&)(_a[j]); float4& s = (float4&)(state[j]); float4 kv; @@ -122,10 +134,10 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons kv.z = k.z * _v; kv.w = k.w * _v; - s.x = s.x * w.x + kv.x + sa * b.x; - s.y = s.y * w.y + kv.y + sa * b.y; - s.z = s.z * w.z + kv.z + sa * b.z; - s.w = s.w * w.w + kv.w + sa * b.w; + s.x = s.x * w.x + kv.x + sa * kk.x * a.x; + s.y = s.y * w.y + kv.y + sa * kk.y * a.y; + s.z = s.z * w.z + kv.z + sa * kk.z * a.z; + s.w = s.w * w.w + kv.w + sa * kk.w * a.w; y += s.x * r.x; y += s.y * r.y; @@ -133,11 +145,12 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons y += s.w * r.w; } dst[t] = y; + dst[T * C + t] = _v * rk; } #pragma unroll for (int i = 0; i < head_size; i++) { - dst[T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; + dst[2 * T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; } } @@ -174,11 +187,12 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const float * w_d = (const float *)dst->src[1]->data; const float * k_d = (const float *)dst->src[2]->data; const float * v_d = (const float *)dst->src[3]->data; - const float * a_d = (const float *)dst->src[4]->data; - const float * b_d = (const float *)dst->src[5]->data; - const float * s_d = (const float *)dst->src[6]->data; + const float * kk_d = (const float *)dst->src[4]->data; + const float * a_d = (const float *)dst->src[5]->data; + const float * r_k_d = (const float *)dst->src[6]->data; + const float * s_d = (const float *)dst->src[7]->data; - const int64_t B = dst->src[6]->ne[1]; + const int64_t B = dst->src[7]->ne[1]; const int64_t T = dst->src[0]->ne[2]; const int64_t C = dst->ne[0]; const int64_t H = dst->src[0]->ne[1]; @@ -187,13 +201,13 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) cudaStream_t stream = ctx.stream(); - GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32); + GGML_ASSERT(dst->src[7]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2); if (C / H == CUDA_WKV_BLOCK_SIZE) { - rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d); + rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); } else { - rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d); + rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); } } diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 0e1f1de4577d..95b717ac4848 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -576,7 +576,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_ } break; case GGML_OP_RWKV_WKV7: { - GGML_ASSERT(op->src[6]->type == GGML_TYPE_F32); + GGML_ASSERT(op->src[7]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == 64); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 18656b346f21..a5d9d746c0aa 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1561,7 +1561,7 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[6]->ne[1]; + const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[7]->ne[1]; const int64_t T = op->src[0]->ne[2]; const int64_t C = op->ne[0]; const int64_t H = op->src[0]->ne[1]; @@ -1579,6 +1579,7 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); if (op->op == GGML_OP_RWKV_WKV7) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), ida++); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[7]), ida++); } ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); ggml_metal_encoder_set_bytes (enc, (void *) &B, sizeof(B), ida++); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 25e78e100898..96521351e85c 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2480,8 +2480,9 @@ kernel void kernel_rwkv_wkv7_f32( device const float * w, device const float * k, device const float * v, + device const float * kk, device const float * a, - device const float * b, + device const float * r_k, device const float * state_in, device float * dst, constant uint & B, @@ -2507,8 +2508,9 @@ kernel void kernel_rwkv_wkv7_f32( threadgroup float _r[head_size]; threadgroup float _w[head_size]; threadgroup float _k[head_size]; + threadgroup float _kk[head_size]; threadgroup float _a[head_size]; - threadgroup float _b[head_size]; + threadgroup float _r_k[head_size]; float state[head_size]; @@ -2516,41 +2518,50 @@ kernel void kernel_rwkv_wkv7_f32( state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + tid * head_size + i]; } + _r_k[tid] = r_k[head_id * head_size + tid]; const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; for (uint t = start_t; t < end_t; t += C) { threadgroup_barrier(mem_flags::mem_threadgroup); + const float w_scale = -0.6065306597126334f; _r[tid] = r[t]; - _w[tid] = w[t]; + _w[tid] = exp(w_scale / (1.0f + exp(-w[t]))); _k[tid] = k[t]; + _kk[tid] = kk[t]; _a[tid] = a[t]; - _b[tid] = b[t]; threadgroup_barrier(mem_flags::mem_threadgroup); const float v_val = v[t]; - float y = 0.0, sa = 0.0; + float y = 0.0, sa = 0.0, rk = 0.0; float4 sa_vec(0.0); + float4 rk_vec(0.0); for (uint j = 0; j < head_size; j += 4) { - float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]); + float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); + float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); + float4 kk_vec = float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + float4 r_k_vec = float4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]); - sa_vec += a_vec * s_vec; + sa_vec += kk_vec * s_vec; + rk_vec += r_vec * k_vec * r_k_vec; } - sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3]; + sa = -(sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3]); + rk = rk_vec[0] + rk_vec[1] + rk_vec[2] + rk_vec[3]; for (uint j = 0; j < head_size; j += 4) { float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]); float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); - float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]); + float4 kk_vec = float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]); float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]); float4 kv = k_vec * v_val; - s_vec = s_vec * w_vec + kv + sa * b_vec; + s_vec = s_vec * w_vec + kv + sa * kk_vec * a_vec; y += dot(s_vec, r_vec); state[j] = s_vec[0]; @@ -2560,10 +2571,11 @@ kernel void kernel_rwkv_wkv7_f32( } dst[t] = y; + dst[T * C + t] = v_val * rk; } for (uint i = 0; i < head_size; i++) { - dst[T * C + batch_id * state_size + head_id * head_size * head_size + dst[2 * T * C + batch_id * state_size + head_id * head_size * head_size + tid * head_size + i] = state[i]; } } diff --git a/ggml/src/ggml-sycl/wkv.cpp b/ggml/src/ggml-sycl/wkv.cpp index b56e0c2400f4..c8a4832e1059 100644 --- a/ggml/src/ggml-sycl/wkv.cpp +++ b/ggml/src/ggml-sycl/wkv.cpp @@ -100,7 +100,7 @@ template static void rwkv_wkv7_f32_kernel( const int B, const int T, const int C, const int H, const float* r, const float* w, const float* k, const float* v, - const float* a, const float* b, const float* s, + const float* kk, const float* a, const float* r_k, const float* s, float* dst, const sycl::nd_item<3>& item_ct1, float* shared_mem) { const int tid = item_ct1.get_local_id(2); @@ -115,8 +115,9 @@ static void rwkv_wkv7_f32_kernel( float* _r = shared_mem; float* _w = _r + head_size; float* _k = _w + head_size; - float* _a = _k + head_size; - float* _b = _a + head_size; + float* _kk = _k + head_size; + float* _a = _kk + head_size; + float* _r_k = _a + head_size; float state[block_size]; @@ -124,6 +125,7 @@ static void rwkv_wkv7_f32_kernel( for (int i = 0; i < head_size; i++) { state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i]; } + _r_k[tid] = r_k[head_i * head_size + tid]; for (int t = batch_i * n_seq_tokens * C + head_i * head_size + tid; t < (batch_i + 1) * n_seq_tokens * C + head_i * head_size + tid; @@ -131,37 +133,44 @@ static void rwkv_wkv7_f32_kernel( item_ct1.barrier(sycl::access::fence_space::local_space); + constexpr float w_scale = -0.6065306597126334f; _r[tid] = r[t]; - _w[tid] = w[t]; + _w[tid] = sycl::exp(w_scale / (1.0f + sycl::exp(-w[t]))); _k[tid] = k[t]; + _kk[tid] = kk[t]; _a[tid] = a[t]; - _b[tid] = b[t]; item_ct1.barrier(sycl::access::fence_space::local_space); const float _v = v[t]; - float y = 0, sa = 0; - sycl::float4 a4, s4; + float y = 0, sa = 0, rk = 0; + sycl::float4 r4, k4, kk4, r_k4, s4; #pragma unroll for (int j = 0; j < head_size; j += 4) { - a4 = sycl::float4(_a[j], _a[j+1], _a[j+2], _a[j+3]); + r4 = sycl::float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); + k4 = sycl::float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); + kk4 = sycl::float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + r_k4 = sycl::float4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); s4 = sycl::float4(state[j], state[j+1], state[j+2], state[j+3]); - sa += sycl::dot(a4, s4); + sa += sycl::dot(kk4, s4); + rk += sycl::dot(r4 * k4, r_k4); } + sa = -sa; - sycl::float4 r4, w4, k4, b4; + sycl::float4 w4, kk4_update, a4; #pragma unroll for (int j = 0; j < head_size; j += 4) { r4 = sycl::float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); w4 = sycl::float4(_w[j], _w[j+1], _w[j+2], _w[j+3]); k4 = sycl::float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); - b4 = sycl::float4(_b[j], _b[j+1], _b[j+2], _b[j+3]); + kk4_update = sycl::float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + a4 = sycl::float4(_a[j], _a[j+1], _a[j+2], _a[j+3]); s4 = sycl::float4(state[j], state[j+1], state[j+2], state[j+3]); sycl::float4 kv4 = k4 * _v; - s4 = s4 * w4 + kv4 + sa * b4; + s4 = s4 * w4 + kv4 + sa * kk4_update * a4; y += sycl::dot(r4, s4); state[j] = s4.x(); @@ -171,11 +180,12 @@ static void rwkv_wkv7_f32_kernel( } dst[t] = y; + dst[T * C + t] = _v * rk; } #pragma unroll for (int i = 0; i < head_size; i++) { - dst[T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; + dst[2 * T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; } } @@ -241,24 +251,25 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { const float* w_d = (const float*)dst->src[1]->data; const float* k_d = (const float*)dst->src[2]->data; const float* v_d = (const float*)dst->src[3]->data; - const float* a_d = (const float*)dst->src[4]->data; - const float* b_d = (const float*)dst->src[5]->data; - const float* s_d = (const float*)dst->src[6]->data; + const float* kk_d = (const float*)dst->src[4]->data; + const float* a_d = (const float*)dst->src[5]->data; + const float* r_k_d = (const float*)dst->src[6]->data; + const float* s_d = (const float*)dst->src[7]->data; float* dst_d = (float*)dst->data; - const int64_t B = dst->src[6]->ne[1]; + const int64_t B = dst->src[7]->ne[1]; const int64_t T = dst->src[0]->ne[2]; const int64_t C = dst->ne[0]; const int64_t H = dst->src[0]->ne[1]; - GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32); + GGML_ASSERT(dst->src[7]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == WKV_BLOCK_SIZE || C / H == WKV_BLOCK_SIZE * 2); dpct::queue_ptr stream = ctx.stream(); // Calculate execution configuration - const size_t shared_mem_size = C / H * 5 * sizeof(float); // For r, w, k, a, b + const size_t shared_mem_size = C / H * 6 * sizeof(float); // For r, w, k, kk, a, r_k sycl::range<3> block_dims(1, 1, C / H); sycl::range<3> grid_dims(1, 1, B * H); @@ -271,7 +282,7 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) { rwkv_wkv7_f32_kernel( - B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d, + B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d, item_ct1, (float*)shared_mem_acc.get_multi_ptr().get() ); }); @@ -284,7 +295,7 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) { rwkv_wkv7_f32_kernel( - B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d, + B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d, item_ct1, (float*)shared_mem_acc.get_multi_ptr().get() ); }); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c0ab9f1c6584..c9e120c0cc13 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5420,7 +5420,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); - ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 9, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; @@ -11929,7 +11929,7 @@ static void ggml_vk_add_id(ggml_backend_vk_context * ctx, vk_context& subctx, co static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, int version) { GGML_ASSERT(version == 6 || version == 7); - int num_srcs = version == 6 ? 6 : 7; + int num_srcs = version == 6 ? 6 : 8; for (int i = 0; i < num_srcs; i++) { GGML_ASSERT(!ggml_is_quantized(dst->src[i]->type)); @@ -11943,7 +11943,7 @@ static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); - vk_subbuffer src_buf[7] = {}; + vk_subbuffer src_buf[8] = {}; for (int i = 0; i < num_srcs; i++) { src_buf[i] = ggml_vk_tensor_subbuffer(ctx, dst->src[i]); } @@ -11960,7 +11960,7 @@ static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx pc, elements); } else if (version == 7) { ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], dst_buf}, + {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], src_buf[7], dst_buf}, pc, elements); } else { // shouldn't happen @@ -11990,7 +11990,7 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx, const size_t seq_length = dst->src[0]->ne[2]; const size_t n_embed = dst->ne[0]; const size_t n_heads = dst->src[0]->ne[1]; - const size_t n_seqs = dst->src[6]->ne[1]; + const size_t n_seqs = dst->src[7]->ne[1]; ggml_vk_op_f32_wkv( ctx, subctx, dst, @@ -18574,7 +18574,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * src_clone[2], src_clone[3], src_clone[4], src_clone[5]); } else if (tensor->op == GGML_OP_RWKV_WKV7) { tensor_clone = ggml_rwkv_wkv7(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], - src_clone[4], src_clone[5], src_clone[6]); + src_clone[4], src_clone[5], src_clone[6], src_clone[7]); } else if (tensor->op == GGML_OP_GATED_DELTA_NET) { tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp index 88c1c02b32b8..851831fd8d61 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp @@ -16,12 +16,13 @@ layout(binding = 0) readonly buffer RBuf { A_TYPE r[]; }; layout(binding = 1) readonly buffer WBuf { A_TYPE w[]; }; layout(binding = 2) readonly buffer KBuf { A_TYPE k[]; }; layout(binding = 3) readonly buffer VBuf { A_TYPE v[]; }; -layout(binding = 4) readonly buffer ABuf { A_TYPE a[]; }; -layout(binding = 5) readonly buffer BBuf { A_TYPE b[]; }; -layout(binding = 6) readonly buffer StateBuf { A_TYPE state_in[]; }; -layout(binding = 7) buffer DstBuf { A_TYPE dst[]; }; +layout(binding = 4) readonly buffer KKBuf { A_TYPE kk[]; }; +layout(binding = 5) readonly buffer ABuf { A_TYPE a[]; }; +layout(binding = 6) readonly buffer RKBuf { A_TYPE r_k[]; }; +layout(binding = 7) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 8) buffer DstBuf { A_TYPE dst[]; }; -shared A_TYPE _r[BLOCK_SIZE], _w[BLOCK_SIZE], _k[BLOCK_SIZE], _a[BLOCK_SIZE], _b[BLOCK_SIZE]; +shared A_TYPE _r[BLOCK_SIZE], _w[BLOCK_SIZE], _k[BLOCK_SIZE], _kk[BLOCK_SIZE], _a[BLOCK_SIZE], _r_k[BLOCK_SIZE]; void main() { const uint head_size = BLOCK_SIZE; @@ -41,25 +42,33 @@ void main() { state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + tid * head_size + i]; } + _r_k[tid] = r_k[head_id * head_size + tid]; const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; for (uint t = start_t; t < end_t; t += C) { barrier(); + const A_TYPE w_scale = A_TYPE(-0.6065306597126334); _r[tid] = r[t]; - _w[tid] = w[t]; + _w[tid] = exp(w_scale / (A_TYPE(1.0) + exp(-w[t]))); _k[tid] = k[t]; + _kk[tid] = kk[t]; _a[tid] = a[t]; - _b[tid] = b[t]; barrier(); A_TYPE sa = 0.0; + A_TYPE rk = 0.0; [[unroll]] for (uint j = 0; j < head_size; j += 4) { + vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); + vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); - vec4 a_vec = vec4(_a[j], _a[j+1], _a[j+2], _a[j+3]); - sa += dot(s_vec, a_vec); + vec4 kk_vec = vec4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + vec4 r_k_vec = vec4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); + sa += dot(s_vec, kk_vec); + rk += dot(r_vec * k_vec, r_k_vec); } + sa = -sa; const A_TYPE v_val = v[t]; A_TYPE y = 0.0; @@ -68,11 +77,12 @@ void main() { vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); vec4 w_vec = vec4(_w[j], _w[j+1], _w[j+2], _w[j+3]); vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); - vec4 b_vec = vec4(_b[j], _b[j+1], _b[j+2], _b[j+3]); + vec4 kk_vec = vec4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); + vec4 a_vec = vec4(_a[j], _a[j+1], _a[j+2], _a[j+3]); vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); vec4 kv = k_vec * v_val; - s_vec = s_vec * w_vec + kv + sa * b_vec; + s_vec = s_vec * w_vec + kv + sa * kk_vec * a_vec; y += dot(r_vec, s_vec); state[j] = s_vec.x; @@ -82,10 +92,11 @@ void main() { } dst[t] = y; + dst[T * C + t] = v_val * rk; } [[unroll]] for (uint i = 0; i < head_size; i++) { - dst[T * C + batch_id * state_size + head_id * head_size * head_size + dst[2 * T * C + batch_id * state_size + head_id * head_size * head_size + tid * head_size + i] = state[i]; } } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 0f682fd1856c..ea3f4b5e84ca 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5812,15 +5812,17 @@ struct ggml_tensor * ggml_rwkv_wkv7( struct ggml_tensor * w, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * kk, struct ggml_tensor * a, - struct ggml_tensor * b, + struct ggml_tensor * r_k, struct ggml_tensor * state) { GGML_ASSERT(ggml_is_contiguous(r)); GGML_ASSERT(ggml_is_contiguous(w)); GGML_ASSERT(ggml_is_contiguous(k)); GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(kk)); GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(b)); + GGML_ASSERT(ggml_is_contiguous(r_k)); GGML_ASSERT(ggml_is_contiguous(state)); const int64_t S = k->ne[0]; @@ -5831,13 +5833,14 @@ struct ggml_tensor * ggml_rwkv_wkv7( GGML_ASSERT(w->ne[0] == S && w->ne[1] == H && w->ne[2] == n_tokens); GGML_ASSERT(k->ne[0] == S && k->ne[1] == H && k->ne[2] == n_tokens); GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(kk->ne[0] == S && kk->ne[1] == H && kk->ne[2] == n_tokens); GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); - GGML_ASSERT(b->ne[0] == S && b->ne[1] == H && b->ne[2] == n_tokens); + GGML_ASSERT(r_k->ne[0] == S && r_k->ne[1] == H); GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); } - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + // concat output, v*rk, and new_state + const int64_t ne[4] = { S * H, 2*n_tokens + S * n_seqs, 1, 1 }; struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); result->op = GGML_OP_RWKV_WKV7; @@ -5845,9 +5848,10 @@ struct ggml_tensor * ggml_rwkv_wkv7( result->src[1] = w; result->src[2] = k; result->src[3] = v; - result->src[4] = a; - result->src[5] = b; - result->src[6] = state; + result->src[4] = kk; + result->src[5] = a; + result->src[6] = r_k; + result->src[7] = state; return result; } diff --git a/src/models/rwkv7-base.cpp b/src/models/rwkv7-base.cpp index 7fcab77745c7..78d38b821452 100644 --- a/src/models/rwkv7-base.cpp +++ b/src/models/rwkv7-base.cpp @@ -67,7 +67,6 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in ggml_tensor * w = ggml_add( ctx0, ggml_mul_mat(ctx0, layer.time_mix_w2, ggml_tanh(ctx0, ggml_mul_mat(ctx0, layer.time_mix_w1, xw))), layer.time_mix_w0); - w = ggml_exp(ctx0, ggml_scale(ctx0, ggml_sigmoid(ctx0, w), -0.606531)); ggml_tensor * k = build_lora_mm(layer.time_mix_key, xk); ggml_tensor * v = build_lora_mm(layer.time_mix_value, xv); @@ -103,10 +102,12 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in a = ggml_reshape_3d(ctx0, a, head_size, head_count, n_tokens); ggml_tensor * wkv_state = build_rs(inp, mctx_cur->get_s_l(il), hparams.n_embd_s(), n_seqs); + ggml_tensor * r_k = ggml_reshape_2d(ctx0, layer.time_mix_r_k, head_size, head_count); - ggml_tensor * wkv_output = ggml_rwkv_wkv7(ctx0, r, w, k, v, ggml_neg(ctx0, kk), ggml_mul(ctx0, kk, a), wkv_state); + ggml_tensor * wkv_output = ggml_rwkv_wkv7(ctx0, r, w, k, v, kk, a, r_k, wkv_state); cur = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, 0); - wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float)); + ggml_tensor * vk = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, n_embd * n_tokens * sizeof(float)); + wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, 2 * n_embd * n_tokens * sizeof(float)); ggml_build_forward_expand( gf, ggml_cpy(ctx0, wkv_state, @@ -124,9 +125,7 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in } else { cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens); } - ggml_tensor * rk = ggml_sum_rows( - ctx0, ggml_mul(ctx0, ggml_mul(ctx0, k, r), ggml_reshape_2d(ctx0, layer.time_mix_r_k, head_size, head_count))); - cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, ggml_mul(ctx0, v, rk), n_embd, n_tokens)); + cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, vk, n_embd, n_tokens)); if (has_gating) { cur = ggml_mul(ctx0, cur, g); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 34811bd50491..436a0a3df65e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4002,13 +4002,14 @@ struct test_rwkv_wkv7 : public test_case { ggml_tensor * w = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); ggml_tensor * k = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); ggml_tensor * v = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); + ggml_tensor * kk = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); ggml_tensor * a = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); + ggml_tensor * r_k = ggml_new_tensor(ctx, type, 2, std::vector{ head_size, head_count }.data()); // Outputs may become NaN with long seqlen without these normalization + kk = ggml_l2_norm(ctx, kk, 1e-7F); a = ggml_l2_norm(ctx, a, 1e-7F); - b = ggml_l2_norm(ctx, b, 1e-7F); ggml_tensor * s = ggml_new_tensor(ctx, type, 2, std::vector{ head_size * head_size * head_count, n_seqs }.data()); - ggml_tensor * out = ggml_rwkv_wkv7(ctx, r, w, k, v, a, b, s); + ggml_tensor * out = ggml_rwkv_wkv7(ctx, r, w, k, v, kk, a, r_k, s); return out; } }; From bf4320f08c6d86f1856f4bd58764bc03e6200895 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 15:15:30 +0800 Subject: [PATCH 02/16] ggml: fuse norm affine operations --- ggml/src/ggml-cpu/ggml-cpu.c | 31 ++++ ggml/src/ggml-cpu/ops.cpp | 94 ++++++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 28 ++- ggml/src/ggml-cuda/norm.cu | 292 +++++++++++++++++++++++++++++++- ggml/src/ggml-cuda/norm.cuh | 7 + 6 files changed, 439 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index eb8341c9aecc..0586c0e1128d 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2993,6 +2993,37 @@ static int ggml_cpu_try_fuse_ops( struct ggml_tensor * node = cgraph->nodes[node_n]; + if (node->op == GGML_OP_NORM) { + // NORM + MUL + ADD fusion + const enum ggml_op fuse_ops[] = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }; + if (ggml_can_fuse(cgraph, node_n, fuse_ops, 3)) { + struct ggml_tensor * mul_node = cgraph->nodes[node_n + 1]; + struct ggml_tensor * add_node = cgraph->nodes[node_n + 2]; + const struct ggml_tensor * mul_w = (mul_node->src[0] == node) + ? mul_node->src[1] : mul_node->src[0]; + const struct ggml_tensor * add_b = (add_node->src[0] == mul_node) + ? add_node->src[1] : add_node->src[0]; + + if (node->src[0]->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && + mul_node->type == GGML_TYPE_F32 && + add_node->type == GGML_TYPE_F32 && + mul_w->type == GGML_TYPE_F32 && + add_b->type == GGML_TYPE_F32 && + mul_w->ne[0] == node->ne[0] && + add_b->ne[0] == node->ne[0] && + ggml_are_same_shape(node, add_node) && + ggml_is_contiguous_rows(node->src[0]) && + ggml_is_contiguous_rows(mul_w) && + ggml_is_contiguous_rows(add_b) && + ggml_is_contiguous_rows(add_node)) { + + ggml_compute_forward_norm_mul_add_fused(params, node, mul_node, add_node); + return 2; + } + } + } + if (node->op == GGML_OP_RMS_NORM) { // RMS_NORM + MUL fusion const enum ggml_op fuse_ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL }; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index afbc98ffc1d6..8bf292eae1e9 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -3768,6 +3768,100 @@ void ggml_compute_forward_norm( } } +// Fused NORM + MUL + ADD: computes dst = norm(src0) * src1 + src2 in a single pass. +static void ggml_compute_forward_norm_mul_add_fused_f32( + const ggml_compute_params * params, + ggml_tensor * dst_norm, + ggml_tensor * dst_mul, + ggml_tensor * dst_add) { + + GGML_ASSERT(dst_mul->src[0] == dst_norm || dst_mul->src[1] == dst_norm); + GGML_ASSERT(dst_add->src[0] == dst_mul || dst_add->src[1] == dst_mul); + + const ggml_tensor * src0 = dst_norm->src[0]; + const ggml_tensor * src1 = dst_mul->src[0] == dst_norm ? dst_mul->src[1] : dst_mul->src[0]; + const ggml_tensor * src2 = dst_add->src[0] == dst_mul ? dst_add->src[1] : dst_add->src[0]; + ggml_tensor * dst = dst_add; + + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + GGML_TENSOR_TERNARY_OP_LOCALS + + float eps; + memcpy(&eps, dst_norm->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + const float * x = (const float *) ((const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + + float sum = 0.0f; + ggml_vec_sum_f32(ne00, &sum, x); + const float mean = sum/ne00; + + float variance = 0.0f; + for (int64_t i00 = 0; i00 < ne00; i00++) { + const float v = x[i00] - mean; + y[i00] = v; + variance += v * v; + } + variance /= ne00; + + const float scale = 1.0f/sqrtf(variance + eps); + + const int64_t i11 = i01 % ne11; + const int64_t i12 = i02 % ne12; + const int64_t i13 = i03 % ne13; + const int64_t i21 = i01 % ne21; + const int64_t i22 = i02 % ne22; + const int64_t i23 = i03 % ne23; + const float * w = (const float *) ((const char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13); + const float * b = (const float *) ((const char *) src2->data + i21*nb21 + i22*nb22 + i23*nb23); + + for (int64_t i00 = 0; i00 < ne00; i00++) { + y[i00] = y[i00] * scale * w[i00] + b[i00]; + } + } + } + } +} + +void ggml_compute_forward_norm_mul_add_fused( + const ggml_compute_params * params, + ggml_tensor * dst_norm, + ggml_tensor * dst_mul, + ggml_tensor * dst_add) { + + GGML_ASSERT(dst_mul != nullptr); + GGML_ASSERT(dst_add != nullptr); + + const ggml_tensor * src0 = dst_norm->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_norm_mul_add_fused_f32(params, dst_norm, dst_mul, dst_add); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_group_rms_norm // fusion kinds that can be combined with the rms_norm computation in a single pass. diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index a8e18c716db7..e707f25f08ca 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -43,6 +43,7 @@ void ggml_compute_forward_repeat_back(const struct ggml_compute_params * params, void ggml_compute_forward_concat(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_silu_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_norm_mul_add_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_norm, struct ggml_tensor * dst_mul, struct ggml_tensor * dst_add); void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rms_norm_mul_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_rms_norm, struct ggml_tensor * dst_mul); void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index dead93f74765..c2d649836709 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3790,8 +3790,10 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } - if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { - const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; + if ((ops.size() == 2 || ops.size() == 3) && + (ops.begin()[0] == GGML_OP_NORM || ops.begin()[0] == GGML_OP_RMS_NORM) && + ops.begin()[1] == GGML_OP_MUL) { + const ggml_tensor *norm = cgraph->nodes[node_idx]; const ggml_tensor *mul = cgraph->nodes[node_idx+1]; const ggml_tensor *add = nullptr; @@ -3799,10 +3801,10 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, add = cgraph->nodes[node_idx+2]; } - GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); + GGML_ASSERT(norm->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(norm->type == GGML_TYPE_F32); - //rms norm only supports F32 + // norm and rms norm only support F32 if (mul->src[0]->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { @@ -3815,12 +3817,12 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } - //if rms norm is the B operand, then we don't handle broadcast - if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { + // If norm is the B operand, then we don't handle broadcast. + if (norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], norm)) { return false; } - //rms_norm kernel assumes contiguous rows + // norm kernels assume contiguous rows if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { return false; } @@ -4356,6 +4358,16 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return fused_node_count - 1; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { + ggml_cuda_op_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + return 2; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_NORM, GGML_OP_MUL }, {})) { + ggml_cuda_op_norm_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); + return 1; + } + 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; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 09d9f3a7d624..af0e73b3d0c4 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1,10 +1,31 @@ #include "norm.cuh" #include -template +template static __global__ void norm_f32( - const float * x, 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 * x, + 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), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0)) { const int nrows = gridDim.x; const int nchannels = gridDim.y; @@ -16,6 +37,22 @@ static __global__ void norm_f32( x += sample*stride_sample + channel*stride_channel + row*stride_row; dst += ((sample*nchannels + channel)*nrows + row)*ncols; + static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); + + 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; + } + + if constexpr (do_add) { + const int add_row = fastmodulo(row, add_nrows_packed); + const int add_channel = fastmodulo(channel, add_nchannels_packed); + const int add_sample = fastmodulo(sample, add_nsamples_packed); + add += add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; + } + float2 mean_var = make_float2(0.0f, 0.0f); ggml_cuda_pdl_sync(); @@ -34,7 +71,17 @@ static __global__ void norm_f32( const float inv_std = rsqrtf(var + eps); for (int col = tid; col < ncols; col += block_size) { - dst[col] = (x[col] - mean) * inv_std; + const float norm = (x[col] - mean) * inv_std; + if constexpr (do_multiply && do_add) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + const int add_col = fastmodulo(col, add_ncols_packed); + dst[col] = norm * mul[mul_col] + add[add_col]; + } else if constexpr (do_multiply) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + dst[col] = norm * mul[mul_col]; + } else { + dst[col] = norm; + } } } @@ -283,10 +330,95 @@ static void norm_f32_cuda( const dim3 blocks_num(nrows, nchannels, nsamples); if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); - norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); } else { const dim3 block_dims(1024, 1, 1); - norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + norm_f32<1024, false><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } +} + +static void norm_mul_f32_cuda(const float * x, + const float * mul, + const float * add, + 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 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 int64_t add_stride_row, + const int64_t add_stride_channel, + const int64_t add_stride_sample, + const uint32_t add_ncols, + const uint32_t add_nrows, + const uint32_t add_nchannels, + const uint32_t add_nsamples, + const float eps, + cudaStream_t stream) { + const dim3 blocks_num(nrows, nchannels, nsamples); + if (mul == nullptr) { + norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream); + return; + } + + 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 (add == nullptr) { + if (ncols < 1024) { + const dim3 block_dims(WARP_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 0, stream}; + ggml_cuda_kernel_launch(norm_f32, launch_params, + x, 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, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } 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(float2): 0, stream}; + ggml_cuda_kernel_launch(norm_f32<1024, true>, launch_params, + x, 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, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } + } else { + const uint3 add_ncols_packed = init_fastdiv_values(add_ncols); + const uint3 add_nrows_packed = init_fastdiv_values(add_nrows); + const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels); + const uint3 add_nsamples_packed = init_fastdiv_values(add_nsamples); + if (ncols < 1024) { + const dim3 block_dims(WARP_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 0, stream}; + ggml_cuda_kernel_launch(norm_f32, launch_params, + x, 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, + add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_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(float2): 0, stream}; + ggml_cuda_kernel_launch(norm_f32<1024, true, true>, launch_params, + x, 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, + add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed); + } } } @@ -456,6 +588,154 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { + const ggml_tensor * norm_src = (ggml_tensor *) dst->src[0]; + float eps = 0.0f; + + memcpy(&eps, dst->op_params, sizeof(float)); + + const float * src0_d = (const float *) norm_src->data; + const float * mul_d = nullptr; + const ggml_tensor * mul_src = nullptr; + + if (mul_tensor->src[0] == dst) { + mul_d = (float *) mul_tensor->src[1]->data; + mul_src = mul_tensor->src[1]; + } else if (mul_tensor->src[1] == dst) { + mul_d = (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(norm_src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(eps >= 0.0f); + + const int64_t ne00 = norm_src->ne[0]; + const int64_t ne01 = norm_src->ne[1]; + const int64_t ne02 = norm_src->ne[2]; + const int64_t ne03 = norm_src->ne[3]; + + const size_t ts0 = ggml_type_size(norm_src->type); + GGML_ASSERT(norm_src->nb[0] == ts0); + const int64_t s01 = norm_src->nb[1] / ts0; + const int64_t s02 = norm_src->nb[2] / ts0; + const int64_t s03 = norm_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]; + + norm_mul_f32_cuda(src0_d, mul_d, nullptr, dst_d, + ne00, ne01, ne02, ne03, + /*s00*/ s01, s02, s03, + /*mul_s00*/ mul_s01, mul_s02, mul_s03, + mul_ncols, mul_nrows, mul_nchannels, mul_nsamples, + /*add_s00*/ 0, 0, 0, + 0, 0, 0, 0, + eps, stream); +} + +void ggml_cuda_op_norm_fused_add(ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + ggml_tensor * mul_tensor, + ggml_tensor * add_tensor) { + const ggml_tensor * norm_src = (ggml_tensor *) dst->src[0]; + float eps = 0.0f; + + memcpy(&eps, dst->op_params, sizeof(float)); + + const float * src0_d = (const float *) norm_src->data; + const float * mul_d = nullptr; + const ggml_tensor * mul_src = nullptr; + + if (mul_tensor->src[0] == dst) { + mul_d = (float *) mul_tensor->src[1]->data; + mul_src = mul_tensor->src[1]; + } else if (mul_tensor->src[1] == dst) { + mul_d = (float *) mul_tensor->src[0]->data; + mul_src = mul_tensor->src[0]; + } else { + GGML_ASSERT(false); + } + + const float * add_d = nullptr; + const ggml_tensor * add_src = nullptr; + + if (add_tensor->src[0] == mul_tensor) { + add_d = (float *) add_tensor->src[1]->data; + add_src = add_tensor->src[1]; + } else if (add_tensor->src[1] == mul_tensor) { + add_d = (float *) add_tensor->src[0]->data; + add_src = add_tensor->src[0]; + } else { + GGML_ASSERT(false); + } + + float * dst_d = (float *) add_tensor->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(norm_src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(add_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(eps >= 0.0f); + + const int64_t ne00 = norm_src->ne[0]; + const int64_t ne01 = norm_src->ne[1]; + const int64_t ne02 = norm_src->ne[2]; + const int64_t ne03 = norm_src->ne[3]; + + const size_t ts0 = ggml_type_size(norm_src->type); + GGML_ASSERT(norm_src->nb[0] == ts0); + const int64_t s01 = norm_src->nb[1] / ts0; + const int64_t s02 = norm_src->nb[2] / ts0; + const int64_t s03 = norm_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]; + + const size_t ts_add = ggml_type_size(add_src->type); + GGML_ASSERT(add_src->nb[0] == ts_add); + const int64_t add_s01 = add_src->nb[1] / ts_add; + const int64_t add_s02 = add_src->nb[2] / ts_add; + const int64_t add_s03 = add_src->nb[3] / ts_add; + + const int add_ncols = add_src->ne[0]; + const int add_nrows = add_src->ne[1]; + const int add_nchannels = add_src->ne[2]; + const int add_nsamples = add_src->ne[3]; + + norm_mul_f32_cuda(src0_d, mul_d, add_d, dst_d, + ne00, ne01, ne02, ne03, + /*s00*/ s01, s02, s03, + /*mul_s00*/ mul_s01, mul_s02, mul_s03, + mul_ncols, mul_nrows, mul_nchannels, mul_nsamples, + /*add_s00*/ add_s01, add_s02, add_s03, + add_ncols, add_nrows, add_nchannels, add_nsamples, + eps, stream); +} + void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *)src0->data; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a74f6376720a..c0ae32e02282 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -2,6 +2,13 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); + +void ggml_cuda_op_norm_fused_add(ggml_backend_cuda_context & ctx, + ggml_tensor * dst, + ggml_tensor * mul_tensor, + ggml_tensor * add_tensor); + void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); From a47b81859a43689e97149bd28058d92095d1ff6b Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 15:30:04 +0800 Subject: [PATCH 03/16] vulkan: fuse norm affine operations --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 98 ++++++++++++++++++- .../vulkan-shaders/norm_mul_add.comp | 96 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 3 files changed, 192 insertions(+), 3 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/norm_mul_add.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c9e120c0cc13..4927122ecde4 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -842,6 +842,7 @@ struct vk_device_struct { vk_pipeline pipeline_set_rows_i32[GGML_TYPE_COUNT]; vk_pipeline pipeline_set_rows_i64[GGML_TYPE_COUNT]; vk_pipeline pipeline_norm_f32; + vk_pipeline pipeline_norm_mul_add_f32; vk_pipeline pipeline_group_norm_f32; vk_pipeline pipeline_rms_norm_f32; vk_pipeline pipeline_rms_norm_mul_f32; @@ -5096,6 +5097,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_nc_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_norm_mul_add_f32, "norm_mul_add_f32", norm_mul_add_f32_len, norm_mul_add_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); @@ -12438,6 +12440,46 @@ static void ggml_vk_norm(ggml_backend_vk_context * ctx, vk_context& subctx, cons ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_NORM, std::move(p)); } +static void ggml_vk_norm_mul_add(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { + ggml_tensor * norm = cgraph->nodes[node_idx]; + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + ggml_tensor * add = cgraph->nodes[node_idx + 2]; + + const ggml_tensor * src0 = norm->src[0]; + const ggml_tensor * src1 = mul->src[0] == norm ? mul->src[1] : mul->src[0]; + const ggml_tensor * src2 = add->src[0] == mul ? add->src[1] : add->src[0]; + ggml_tensor * dst = add; + + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + float * op_params = (float *)norm->op_params; + + vk_op_binary_push_constants pc { + (uint32_t)ggml_nelements(src0), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2], (uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + op_params[0], 0.0f, 0, + }; + + init_pushconst_fastdiv(pc); + init_pushconst_tensor_offsets(ctx, pc, src0, src1, nullptr, nullptr, dst); + + vk_pipeline pipeline = ctx->device->pipeline_norm_mul_add_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0, true); + vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1, true); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + vk_subbuffer src2_buf = ggml_vk_tensor_subbuffer(ctx, src2); + + std::array elements = { (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3] }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, dst_buf, src2_buf }, pc, elements); +} + static void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { const int * int_op_params = (const int *)dst->op_params; const float * float_op_params = (const float *)dst->op_params; @@ -14753,7 +14795,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_NORM: - ggml_vk_norm(ctx, compute_ctx, src0, node); + if (ctx->num_additional_fused_ops == 2) { + ggml_vk_norm_mul_add(ctx, compute_ctx, cgraph, node_idx); + } else { + ggml_vk_norm(ctx, compute_ctx, src0, node); + } break; case GGML_OP_GROUP_NORM: @@ -15709,7 +15755,41 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g return false; } - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_NORM && ops.begin()[1] == GGML_OP_MUL && ops.begin()[2] == GGML_OP_ADD) { + const ggml_tensor *norm = cgraph->nodes[node_idx]; + const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; + const ggml_tensor *add = cgraph->nodes[node_idx + 2]; + const ggml_tensor *w = mul->src[0] == norm ? mul->src[1] : mul->src[0]; + const ggml_tensor *b = add->src[0] == mul ? add->src[1] : add->src[0]; + + if (norm->src[0]->type != GGML_TYPE_F32 || + norm->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32 || + w->type != GGML_TYPE_F32 || + b->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_are_same_shape(w, b) || + !ggml_are_same_stride(w, b) || + w->ne[0] != norm->ne[0] || + !ggml_are_same_shape(norm, add)) { + return false; + } + + if (!ggml_is_contiguous_rows(norm->src[0]) || + !ggml_is_contiguous_rows(w) || + !ggml_is_contiguous_rows(b) || + !ggml_is_contiguous_rows(add)) { + return false; + } + + // The shader reuses the scale tensor's shape/stride for the bias tensor. + if (get_misalign_bytes(ctx, b) != 0) { + return false; + } + } else if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { // additional constraints specific to this fusion const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; @@ -16362,6 +16442,12 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { + ctx->num_additional_fused_ops = 2; + fusion_string = "NORM_MUL_ADD"; + op_srcs_fused_elementwise[0] = true; + op_srcs_fused_elementwise[1] = true; + op_srcs_fused_elementwise[2] = true; } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 4 }) && ggml_check_edges(cgraph, i, rms_norm_mul_rope_view_set_rows_edges) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i) && @@ -16710,6 +16796,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * if (keep_pattern(snake_pattern)) { continue; } + if (keep_pattern({ GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { + continue; + } // First, grab the next unused node. current_set.push_back(first_unused); @@ -16733,13 +16822,16 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * match_pattern(topk_moe_sigmoid_norm_bias, j) || match_pattern(topk_moe_early_softmax, j) || match_pattern(topk_moe_late_softmax, j) || - match_pattern(snake_pattern, j)) { + match_pattern(snake_pattern, j) || + match_pattern({ GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }, j)) { continue; } bool ok = true; for (int c = first_unused; c < j; ++c) { if (!used[c] && is_src_of(graph->nodes[j], graph->nodes[c]) && + !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_NORM && graph->nodes[j]->op == GGML_OP_MUL) && + !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_MUL && graph->nodes[j]->op == GGML_OP_ADD) && !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_RMS_NORM && graph->nodes[j]->op == GGML_OP_MUL) && !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_MUL_MAT && graph->nodes[j]->op == GGML_OP_ADD) && !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_MUL_MAT_ID && graph->nodes[j]->op == GGML_OP_ADD_ID) && diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/norm_mul_add.comp b/ggml/src/ggml-vulkan/vulkan-shaders/norm_mul_add.comp new file mode 100644 index 000000000000..0d262641a89d --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/norm_mul_add.comp @@ -0,0 +1,96 @@ +#version 450 + +#include "generic_binary_head.glsl" +#include "types.glsl" + +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 + +layout (binding = 3) readonly buffer C {B_TYPE data_c[];}; + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +shared vec2 sum[BLOCK_SIZE]; + +void norm_mul_add(uint num_iters) { + const uint ncols = p.ne00; + const uint row = gl_WorkGroupID.x; + const uint channel = gl_WorkGroupID.y; + const uint samp = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationID.x; + + const uint a_base = get_aoffset() + src0_idx(0, row, channel, samp); + const uint b_base = get_boffset() + src1_idx(0, row, channel, samp); + const uint d_base = get_doffset() + dst_idx(0, row, channel, samp); + + vec2 sum_tid = vec2(0.0f, 0.0f); + + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col < ncols) { + const float xi = float(data_a[a_base + col*p.nb00]); + sum_tid.x += xi; + sum_tid.y += xi * xi; + } + } + + sum[tid] = sum_tid; + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum[tid] += sum[tid + s]; + } + barrier(); + } + + const float mean = sum[0].x / float(ncols); + const float var = sum[0].y / float(ncols) - mean * mean; + const float inv_std = inversesqrt(var + p.param1); + + if (ncols > p.ne10) { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + const uint b_col = fastmod(col, p.ne10); + const float xi = float(data_a[a_base + col*p.nb00]); + const float w = float(data_b[b_base + b_col*p.nb10]); + const float b = float(data_c[b_base + b_col*p.nb10]); + data_d[d_base + col*p.nb20] = D_TYPE((xi - mean) * inv_std * w + b); + } + } else { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + const float xi = float(data_a[a_base + col*p.nb00]); + const float w = float(data_b[b_base + col*p.nb10]); + const float b = float(data_c[b_base + col*p.nb10]); + data_d[d_base + col*p.nb20] = D_TYPE((xi - mean) * inv_std * w + b); + } + } +} + +void main() { + uint num_blocks = (p.ne00 + BLOCK_SIZE - 1) / BLOCK_SIZE; + if (num_blocks > 32) { + norm_mul_add(num_blocks); + } else if (num_blocks > 16) { + norm_mul_add(32); + } else if (num_blocks > 12) { + norm_mul_add(16); + } else if (num_blocks > 10) { + norm_mul_add(12); + } else if (num_blocks > 8) { + norm_mul_add(10); + } else if (num_blocks > 4) { + norm_mul_add(8); + } else if (num_blocks == 4) { + norm_mul_add(4); + } else if (num_blocks == 3) { + norm_mul_add(3); + } else if (num_blocks == 2) { + norm_mul_add(2); + } else if (num_blocks == 1) { + norm_mul_add(1); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 1925582ffedb..33ed69548c51 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -772,6 +772,7 @@ void process_shaders() { // Norms string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("norm_mul_add_f32", "norm_mul_add.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); From c88673f9cacc6b69e79cca9bd5edf10da8456b2e Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 17:07:42 +0800 Subject: [PATCH 04/16] ggml: fuse add-mul elementwise operations --- ggml/src/ggml-cpu/ggml-cpu.c | 71 +++++- ggml/src/ggml-cpu/ops.cpp | 42 ++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml-cuda/binbcast.cu | 77 ++++++ ggml/src/ggml-cuda/binbcast.cuh | 7 + ggml/src/ggml-cuda/ggml-cuda.cu | 131 +++++++++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 223 ++++++++++++++++-- .../ggml-vulkan/vulkan-shaders/add_mul.comp | 33 +++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 9 files changed, 549 insertions(+), 37 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/add_mul.comp diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 0586c0e1128d..3ec719a09aea 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2981,11 +2981,30 @@ struct ggml_cplan ggml_graph_plan( // Returns the number of nodes skipped by fusion (>=1), or 0 if no fusion was applied. static bool ggml_cpu_disable_fusion = false; // initialized once in ggml_cpu_init(), read-only afterwards +static int ggml_cpu_find_single_consumer(const struct ggml_cgraph * cgraph, const int node_n) { + if (ggml_node_get_use_count(cgraph, node_n) != 1) { + return -1; + } + + const struct ggml_tensor * node = cgraph->nodes[node_n]; + for (int i = node_n + 1; i < cgraph->n_nodes; ++i) { + const struct ggml_tensor * consumer = cgraph->nodes[i]; + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + if (consumer->src[src_idx] == node) { + return i; + } + } + } + + return -1; +} + static int ggml_cpu_try_fuse_ops( const struct ggml_cgraph * cgraph, const int node_n, const struct ggml_compute_params * params, - const struct ggml_cplan * cplan) { + const struct ggml_cplan * cplan, + int * skip_node_n) { if (ggml_cpu_disable_fusion || cplan->use_ref) { return 0; @@ -2993,6 +3012,37 @@ static int ggml_cpu_try_fuse_ops( struct ggml_tensor * node = cgraph->nodes[node_n]; + if (node->op == GGML_OP_ADD) { + // ADD + MUL fusion + const int mul_node_n = ggml_cpu_find_single_consumer(cgraph, node_n); + if (mul_node_n > node_n && cgraph->nodes[mul_node_n]->op == GGML_OP_MUL) { + struct ggml_tensor * mul_node = cgraph->nodes[mul_node_n]; + const struct ggml_tensor * scale = (mul_node->src[0] == node) + ? mul_node->src[1] : mul_node->src[0]; + + if (node->src[0]->type == GGML_TYPE_F32 && + node->src[1]->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && + scale->type == GGML_TYPE_F32 && + mul_node->type == GGML_TYPE_F32 && + ggml_are_same_shape(node->src[0], mul_node) && + ggml_are_same_shape(node->src[1], mul_node) && + ggml_are_same_shape(scale, mul_node) && + ggml_is_contiguous(node->src[0]) && + ggml_is_contiguous(node->src[1]) && + ggml_is_contiguous(scale) && + ggml_is_contiguous(mul_node)) { + + ggml_compute_forward_add_mul_fused(params, node, mul_node); + if (mul_node_n == node_n + 1) { + return 1; + } + *skip_node_n = mul_node_n; + return -1; + } + } + } + if (node->op == GGML_OP_NORM) { // NORM + MUL + ADD fusion const enum ggml_op fuse_ops[] = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }; @@ -3074,6 +3124,7 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { GGML_PRINT_DEBUG("thread #%d compute-start cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); #endif + int skip_node_n = -1; for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) { struct ggml_tensor * node = cgraph->nodes[node_n]; @@ -3086,13 +3137,19 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { continue; } - // TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time - // Try fused ops, fall back to normal compute - const int n_fused = ggml_cpu_try_fuse_ops(cgraph, node_n, ¶ms, cplan); - if (n_fused > 0) { - node_n += n_fused; + if (skip_node_n == node_n) { + skip_node_n = -1; } else { - ggml_compute_forward(¶ms, node); + // TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time + // Try fused ops, fall back to normal compute + const int n_fused = ggml_cpu_try_fuse_ops(cgraph, node_n, ¶ms, cplan, &skip_node_n); + if (n_fused != 0) { + if (n_fused > 0) { + node_n += n_fused; + } + } else { + ggml_compute_forward(¶ms, node); + } } if (state->ith == 0 && cplan->abort_callback && diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 8bf292eae1e9..4e41c2a1f93a 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -699,6 +699,48 @@ void ggml_compute_forward_add( } } +void ggml_compute_forward_add_mul_fused( + const ggml_compute_params * params, + ggml_tensor * dst_add, + ggml_tensor * dst_mul) { + GGML_ASSERT(dst_add->op == GGML_OP_ADD); + GGML_ASSERT(dst_mul->op == GGML_OP_MUL); + GGML_ASSERT(dst_mul->src[0] == dst_add || dst_mul->src[1] == dst_add); + + const ggml_tensor * src0 = dst_add->src[0]; + const ggml_tensor * src1 = dst_add->src[1]; + const ggml_tensor * scale = dst_mul->src[0] == dst_add ? dst_mul->src[1] : dst_mul->src[0]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(dst_mul->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(dst_mul)); + GGML_ASSERT(ggml_are_same_shape(src0, dst_mul)); + GGML_ASSERT(ggml_are_same_shape(src1, dst_mul)); + GGML_ASSERT(ggml_are_same_shape(scale, dst_mul)); + + const int64_t n = ggml_nelements(dst_mul); + const int64_t dr = (n + params->nth - 1) / params->nth; + const int64_t i0 = dr * params->ith; + const int64_t i1 = MIN(i0 + dr, n); + + const float * a = (const float *) src0->data; + const float * b = (const float *) src1->data; + const float * s = (const float *) scale->data; + float * dst = (float *) dst_mul->data; + + for (int64_t i = i0; i < i1; ++i) { + dst[i] = (a[i] + b[i]) * s[i]; + } + + GGML_UNUSED(dst_add); +} + // ggml_compute_forward_add_id static void ggml_compute_forward_add_id_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index e707f25f08ca..efc8004ce551 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -29,6 +29,7 @@ extern "C" { void ggml_compute_forward_dup(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_add(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add_mul_fused(const struct ggml_compute_params * params, struct ggml_tensor * dst_add, struct ggml_tensor * dst_mul); void ggml_compute_forward_add_id(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_add1(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_acc(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 6e11dc6e71af..9567de4b4a0d 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -106,6 +106,39 @@ static __global__ void k_mul_sub_add_fused_f32( dst[idx] = base[ib] + (scale[is] * v - v); } +static __global__ void k_add_mul_fused_f32( + const float * __restrict__ src0, + const float * __restrict__ src1, + const float * __restrict__ scale, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s0_ne0, const int64_t s0_ne1, const int64_t s0_ne2, const int64_t s0_ne3, + const int64_t s0_s0, const int64_t s0_s1, const int64_t s0_s2, const int64_t s0_s3, + const int64_t s1_ne0, const int64_t s1_ne1, const int64_t s1_ne2, const int64_t s1_ne3, + const int64_t s1_s0, const int64_t s1_s1, const int64_t s1_s2, const int64_t s1_s3, + const int64_t sc_ne0, const int64_t sc_ne1, const int64_t sc_ne2, const int64_t sc_ne3, + const int64_t sc_s0, const int64_t sc_s1, const int64_t sc_s2, const int64_t sc_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t is0 = ggml_cuda_index_4d(i0, i1, i2, i3, s0_ne0, s0_ne1, s0_ne2, s0_ne3, s0_s0, s0_s1, s0_s2, s0_s3); + const size_t is1 = ggml_cuda_index_4d(i0, i1, i2, i3, s1_ne0, s1_ne1, s1_ne2, s1_ne3, s1_s0, s1_s1, s1_s2, s1_s3); + const size_t isc = ggml_cuda_index_4d(i0, i1, i2, i3, sc_ne0, sc_ne1, sc_ne2, sc_ne3, sc_s0, sc_s1, sc_s2, sc_s3); + + dst[idx] = (src0[is0] + src1[is1]) * scale[isc]; +} + template type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(src0, dst)); + GGML_ASSERT(ggml_can_repeat(src1, dst)); + GGML_ASSERT(ggml_can_repeat(scale, dst)); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_add_mul_fused_f32, launch_params, + (const float *) src0->data, + (const float *) src1->data, + (const float *) scale->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], + stride(src0, 0), stride(src0, 1), stride(src0, 2), stride(src0, 3), + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + stride(src1, 0), stride(src1, 1), stride(src1, 2), stride(src1, 3), + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], + stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3)); +} + void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index b570401cc286..40ec517f2e1f 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -24,3 +24,10 @@ void ggml_cuda_op_mul_sub_add_fused( const ggml_tensor * scale, const ggml_tensor * value, ggml_tensor * dst); + +void ggml_cuda_op_add_mul_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * scale, + ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c2d649836709..e8aa0673f695 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -85,6 +85,8 @@ #include #include #include +#include +#include #include static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); @@ -3663,6 +3665,91 @@ static bool ggml_cuda_should_fuse_mul_sub_add( return true; } +static bool ggml_cuda_should_fuse_add_mul( + const ggml_tensor * add, + const ggml_tensor * mul, + const ggml_tensor ** src0, + const ggml_tensor ** src1, + const ggml_tensor ** scale) { + if (add->op != GGML_OP_ADD || mul->op != GGML_OP_MUL) { + return false; + } + + const ggml_tensor * s = nullptr; + if (mul->src[0] == add) { + s = mul->src[1]; + } else if (mul->src[1] == add) { + s = mul->src[0]; + } else { + return false; + } + + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + + if (a->type != GGML_TYPE_F32 || b->type != GGML_TYPE_F32 || s->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_can_repeat(a, mul) || !ggml_can_repeat(b, mul) || !ggml_can_repeat(s, mul)) { + return false; + } + + if (!ggml_is_contiguous(a) || !ggml_is_contiguous(b) || !ggml_is_contiguous(s) || !ggml_is_contiguous(mul)) { + return false; + } + + *src0 = a; + *src1 = b; + *scale = s; + return true; +} + +static int ggml_cuda_find_single_consumer(const ggml_cgraph * cgraph, int node_idx) { + if (ggml_node_get_use_count(cgraph, node_idx) != 1) { + return -1; + } + + const ggml_tensor * node = cgraph->nodes[node_idx]; + for (int i = node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * consumer = cgraph->nodes[i]; + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + if (consumer->src[src_idx] == node) { + return i; + } + } + } + + return -1; +} + +static bool ggml_cuda_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b) { + if (a == nullptr || b == nullptr || a->buffer == nullptr || b->buffer == nullptr) { + return false; + } + + const uintptr_t a_start = reinterpret_cast(a->data); + const uintptr_t b_start = reinterpret_cast(b->data); + const uintptr_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + const uintptr_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + + return a_start < b_end && b_start < a_end; +} + +static bool ggml_cuda_add_mul_memory_ok( + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * scale, + const ggml_tensor * dst) { + const auto can_read_write_inplace = [dst](const ggml_tensor * src) { + return !ggml_cuda_tensors_overlap(dst, src) || ggml_are_same_layout(dst, src); + }; + + return can_read_write_inplace(src0) && + can_read_write_inplace(src1) && + can_read_write_inplace(scale); +} + // returns whether the write (out) nodes overwrite the read nodes in operation static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, @@ -3950,8 +4037,14 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } -// try and fuse nodes and return the number of nodes to skip -static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { +// Try and fuse nodes. A positive return value is the number of following +// contiguous nodes to skip; -1 means the current node was handled but no +// following nodes can be skipped because the fused consumer is non-contiguous. +static int ggml_cuda_try_fuse( + ggml_backend_cuda_context * cuda_ctx, + ggml_cgraph * cgraph, + int i, + std::unordered_set & fused_noncontiguous_nodes) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { @@ -4087,6 +4180,28 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // Elementwise gate: (a + b) * scale. RWKV builds this as ADD consumed by + // a later MUL, with unrelated view nodes often scheduled between them. + if (node->op == GGML_OP_ADD) { + const int mul_idx = ggml_cuda_find_single_consumer(cgraph, i); + const ggml_tensor * src0 = nullptr; + const ggml_tensor * src1 = nullptr; + const ggml_tensor * scale = nullptr; + + const bool should_fuse = mul_idx > i && cgraph->nodes[mul_idx]->op == GGML_OP_MUL && + ggml_cuda_should_fuse_add_mul(cgraph->nodes[i], cgraph->nodes[mul_idx], &src0, &src1, &scale); + const bool memory_ok = should_fuse && ggml_cuda_add_mul_memory_ok(src0, src1, scale, cgraph->nodes[mul_idx]); + + if (memory_ok) { + ggml_cuda_op_add_mul_fused(*cuda_ctx, src0, src1, scale, cgraph->nodes[mul_idx]); + if (mul_idx == i + 1) { + return 1; + } + fused_noncontiguous_nodes.insert(cgraph->nodes[mul_idx]); + return -1; + } + } + // Snake activation: y = x + sin(a*x)^2 * inv_b // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add if (ggml_can_fuse_subgraph(cgraph, i, @@ -4506,6 +4621,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud stream_ctx.concurrent_events.clear(); } + std::unordered_set fused_noncontiguous_nodes; + for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -4554,10 +4671,16 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); + if (fused_noncontiguous_nodes.erase(node) > 0) { + continue; + } + + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, fused_noncontiguous_nodes); if (nodes_to_skip != 0) { - i += nodes_to_skip; + if (nodes_to_skip > 0) { + i += nodes_to_skip; + } continue; } #ifndef NDEBUG diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 4927122ecde4..cc02f66fa58a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -62,6 +62,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #include #include #include +#include #include #include #include @@ -812,6 +813,7 @@ struct vk_device_struct { vk_pipeline pipeline_sub_norepeat[2][2][2]; vk_pipeline pipeline_mul[2][2][2]; vk_pipeline pipeline_mul_norepeat[2][2][2]; + vk_pipeline pipeline_add_mul_f32; vk_pipeline pipeline_div[2][2][2]; vk_pipeline pipeline_div_norepeat[2][2][2]; vk_pipeline pipeline_add_rms[2][2][2]; @@ -2149,6 +2151,8 @@ struct ggml_backend_vk_context { // Bit 'i' means nodes[start_of_fusion + i] writes to memory. // If there's no fusion, bit 0 is still set. int fused_ops_write_mask {}; + bool fused_add_mul {}; + int fused_add_mul_dst_idx {-1}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; @@ -5074,6 +5078,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows_back_f32, "get_rows_back_f32", get_rows_back_f32_len, get_rows_back_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {256, 1, 1}, {}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_add_mul_f32, "add_mul_f32", add_mul_f32_len, add_mul_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true); for (auto &it : device->pipeline_fa_mask_opt) { @@ -11899,6 +11904,50 @@ static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const }); } +static void ggml_vk_add_mul(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { + const ggml_tensor * add = cgraph->nodes[node_idx]; + const int dst_idx = ctx->fused_add_mul_dst_idx >= 0 ? ctx->fused_add_mul_dst_idx : node_idx + 1; + ggml_tensor * dst = cgraph->nodes[dst_idx]; + const ggml_tensor * scale = dst->src[0] == add ? dst->src[1] : dst->src[0]; + const ggml_tensor * src0 = add->src[0]; + const ggml_tensor * src1 = add->src[1]; + + const uint32_t src0_type_size = ggml_type_size(src0->type); + const uint32_t src1_type_size = ggml_type_size(src1->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + vk_op_binary_push_constants pc { + (uint32_t)ggml_nelements(dst), + (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, + (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2], (uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, + (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2], (uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, + 0, + 0.0f, 0.0f, 0, + }; + + init_pushconst_fastdiv(pc); + init_pushconst_tensor_offsets(ctx, pc, src0, src1, nullptr, nullptr, dst); + + vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0, true); + vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1, true); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + vk_subbuffer scale_buf = ggml_vk_tensor_subbuffer(ctx, scale, true); + + std::array elements; + const uint32_t ne = (uint32_t) ggml_nelements(dst); + if (ne > 262144) { + elements = { 512, 512, CEIL_DIV(ne, 262144) }; + } else if (ne > 512) { + elements = { 512, CEIL_DIV(ne, 512), 1 }; + } else { + elements = { ne, 1, 1 }; + } + + vk_pipeline pipeline = ctx->device->pipeline_add_mul_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, dst_buf, scale_buf }, pc, elements); +} + static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); @@ -14601,24 +14650,41 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr return false; }; - // For all fused ops, check if the destination node or any of the source - // nodes require synchronization. - for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1 && !need_sync; ++i) { - const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; - // If the node actually writes to memory, then check if it needs to sync - if (ctx->fused_ops_write_mask & (1 << i)) { - if (overlaps_unsynced(cur_node, ctx->unsynced_nodes_read) || overlaps_unsynced(cur_node, ctx->unsynced_nodes_written)) { + if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { + const ggml_tensor * add = node; + const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; + const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; + const ggml_tensor * fused_srcs[] = { add->src[0], add->src[1], scale }; + + if (overlaps_unsynced(mul, ctx->unsynced_nodes_read) || overlaps_unsynced(mul, ctx->unsynced_nodes_written)) { + need_sync = true; + } + for (const ggml_tensor * src : fused_srcs) { + if (src && overlaps_unsynced(src, ctx->unsynced_nodes_written)) { need_sync = true; break; } } - for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { - if (!cur_node->src[j]) { - continue; + } else { + // For all fused ops, check if the destination node or any of the source + // nodes require synchronization. + for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1 && !need_sync; ++i) { + const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; + // If the node actually writes to memory, then check if it needs to sync + if (ctx->fused_ops_write_mask & (1 << i)) { + if (overlaps_unsynced(cur_node, ctx->unsynced_nodes_read) || overlaps_unsynced(cur_node, ctx->unsynced_nodes_written)) { + need_sync = true; + break; + } } - if (overlaps_unsynced(cur_node->src[j], ctx->unsynced_nodes_written)) { - need_sync = true; - break; + for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { + if (!cur_node->src[j]) { + continue; + } + if (overlaps_unsynced(cur_node->src[j], ctx->unsynced_nodes_written)) { + need_sync = true; + break; + } } } } @@ -14637,18 +14703,28 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr ggml_vk_sync_buffers(ctx, compute_ctx); } } - // Add all fused nodes to the unsynchronized lists. - for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { - const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; - // Multiple outputs could be written, e.g. in topk_moe. Add them all to the list. - if (ctx->fused_ops_write_mask & (1 << i)) { - ctx->unsynced_nodes_written.push_back(cur_node); - } - for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { - if (!cur_node->src[j]) { - continue; + if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { + const ggml_tensor * add = node; + const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; + const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; + ctx->unsynced_nodes_written.push_back(mul); + ctx->unsynced_nodes_read.push_back(add->src[0]); + ctx->unsynced_nodes_read.push_back(add->src[1]); + ctx->unsynced_nodes_read.push_back(scale); + } else { + // Add all fused nodes to the unsynchronized lists. + for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { + const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; + // Multiple outputs could be written, e.g. in topk_moe. Add them all to the list. + if (ctx->fused_ops_write_mask & (1 << i)) { + ctx->unsynced_nodes_written.push_back(cur_node); + } + for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { + if (!cur_node->src[j]) { + continue; + } + ctx->unsynced_nodes_read.push_back(cur_node->src[j]); } - ctx->unsynced_nodes_read.push_back(cur_node->src[j]); } } } @@ -14690,8 +14766,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_ADD: - if (ctx->num_additional_fused_ops) { - ggml_vk_multi_add(ctx, compute_ctx, cgraph, node_idx); + if (ctx->fused_add_mul) { + ggml_vk_add_mul(ctx, compute_ctx, cgraph, node_idx); + } else if (ctx->num_additional_fused_ops) { + ggml_vk_multi_add(ctx, compute_ctx, cgraph, node_idx); } else { ggml_vk_add(ctx, compute_ctx, src0, src1, node); } @@ -15750,6 +15828,39 @@ static bool ggml_vk_is_empty(ggml_tensor * node) { return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; } +static bool ggml_vk_should_fuse_add_mul(const ggml_backend_vk_context * ctx, const ggml_tensor * add, const ggml_tensor * mul) { + if (add->op != GGML_OP_ADD || mul->op != GGML_OP_MUL) { + return false; + } + if (mul->src[0] != add && mul->src[1] != add) { + return false; + } + + const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; + + if (add->src[0]->type != GGML_TYPE_F32 || + add->src[1]->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32 || + scale->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_are_same_shape(add->src[0], mul) || + !ggml_are_same_shape(add->src[1], mul) || + !ggml_are_same_shape(scale, mul) || + !ggml_are_same_stride(scale, mul)) { + return false; + } + + // The shader reuses the destination shape/stride and offset for scale. + if (get_misalign_bytes(ctx, scale) != 0 || get_misalign_bytes(ctx, mul) != 0) { + return false; + } + + return true; +} + static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; @@ -15789,6 +15900,10 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g if (get_misalign_bytes(ctx, b) != 0) { return false; } + } else if (ops.size() == 2 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_MUL) { + if (!ggml_vk_should_fuse_add_mul(ctx, cgraph->nodes[node_idx], cgraph->nodes[node_idx + 1])) { + return false; + } } else if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { // additional constraints specific to this fusion const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; @@ -16209,6 +16324,32 @@ static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b return false; } +static int ggml_vk_find_single_consumer(const ggml_cgraph * cgraph, int node_idx) { + if (ggml_node_get_use_count(cgraph, node_idx) != 1) { + return -1; + } + + const ggml_tensor * node = cgraph->nodes[node_idx]; + for (int i = node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * consumer = cgraph->nodes[i]; + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + if (consumer->src[src_idx] == node) { + return i; + } + } + } + + return -1; +} + +static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor * mul) { + const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; + + return !ggml_vk_tensors_overlap(add->src[0], mul, true) && + !ggml_vk_tensors_overlap(add->src[1], mul, true) && + !ggml_vk_tensors_overlap(scale, mul, true); +} + static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { GGML_UNUSED(ctx); @@ -16389,11 +16530,16 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg uint64_t batch_flops = 0; uint64_t total_flops = 0; uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u); + std::unordered_set fused_noncontiguous_nodes; for (int i = 0; i < cgraph->n_nodes; i++) { if (first_node_in_batch) { submit_node_idx = i; } + if (fused_noncontiguous_nodes.erase(cgraph->nodes[i]) > 0) { + continue; + } + { auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]); batch_flops += node_flops; @@ -16408,6 +16554,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; + ctx->fused_add_mul = false; + ctx->fused_add_mul_dst_idx = -1; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -16442,6 +16590,22 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_ADD, GGML_OP_MUL })) { + ctx->num_additional_fused_ops = 1; + ctx->fused_add_mul = true; + fusion_string = "ADD_MUL"; + op_srcs_fused_elementwise[0] = true; + op_srcs_fused_elementwise[1] = true; + } else if (cgraph->nodes[i]->op == GGML_OP_ADD) { + const int mul_idx = ggml_vk_find_single_consumer(cgraph, i); + if (mul_idx > i && + cgraph->nodes[mul_idx]->op == GGML_OP_MUL && + ggml_vk_should_fuse_add_mul(ctx, cgraph->nodes[i], cgraph->nodes[mul_idx]) && + ggml_vk_add_mul_memory_ok(cgraph->nodes[i], cgraph->nodes[mul_idx])) { + ctx->fused_add_mul = true; + ctx->fused_add_mul_dst_idx = mul_idx; + fusion_string = "ADD_MUL"; + } } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { ctx->num_additional_fused_ops = 2; fusion_string = "NORM_MUL_ADD"; @@ -16610,6 +16774,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_ops_write_mask = 1; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; + ctx->fused_add_mul = false; + ctx->fused_add_mul_dst_idx = -1; } } @@ -16638,6 +16804,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } if (enqueued) { + if (ctx->fused_add_mul_dst_idx >= 0) { + fused_noncontiguous_nodes.insert(cgraph->nodes[ctx->fused_add_mul_dst_idx]); + } ++submitted_nodes; #ifndef GGML_VULKAN_CHECK_RESULTS @@ -16659,6 +16828,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg i += ctx->num_additional_fused_ops; ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 0; + ctx->fused_add_mul = false; + ctx->fused_add_mul_dst_idx = -1; } ctx->last_total_flops = total_flops; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/add_mul.comp b/ggml/src/ggml-vulkan/vulkan-shaders/add_mul.comp new file mode 100644 index 000000000000..a4187bf49c27 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/add_mul.comp @@ -0,0 +1,33 @@ +#version 450 + +#include "types.glsl" +#include "generic_binary_head.glsl" + +const uint num_threads = 256; + +layout (binding = 3) readonly buffer C {B_TYPE data_c[];}; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + uint idx = get_idx(); + + // num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation. + const uint num_iter = 2; + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (idx >= p.ne) { + continue; + } + uint i00, i01, i02, i03; + get_indices(idx, i00, i01, i02, i03); + + const uint d_idx = get_doffset() + dst_idx(i00, i01, i02, i03); + const FLOAT_TYPE sum = + FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)]); + data_d[d_idx] = D_TYPE(sum * FLOAT_TYPE(data_c[d_idx])); + + idx += num_threads; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 33ed69548c51..b2769b9ca41b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -837,6 +837,7 @@ void process_shaders() { string_to_spv("sub_f32", "sub.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + string_to_spv("add_mul_f32", "add_mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); From 6bf3c2b08e8e211312fb7cedfe28fc50c7d08876 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 17:51:43 +0800 Subject: [PATCH 05/16] cuda: specialize rwkv lerp fusion --- ggml/src/ggml-cuda/binbcast.cu | 46 ++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 9567de4b4a0d..fa0d03ec55e7 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -72,6 +72,24 @@ static __global__ void k_lerp_fused_f32( dst[idx] = c + (x_prev[ixp] - c) * weight[iw]; } +static __global__ void k_lerp_fused_rwv_contig_f32( + const float * __restrict__ x_prev, + const float * __restrict__ cur, + const float * __restrict__ weight, + float * __restrict__ dst, + const int64_t ne0, + const int64_t base_total) { + const int64_t ibase = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (ibase >= base_total) { + return; + } + + const int64_t imix = blockIdx.y; + const int64_t i0 = ibase % ne0; + const float c = cur[ibase]; + dst[imix * base_total + ibase] = c + (x_prev[ibase] - c) * weight[imix * ne0 + i0]; +} + static __global__ void k_mul_sub_add_fused_f32( const float * __restrict__ base, const float * __restrict__ scale, @@ -680,6 +698,34 @@ void ggml_cuda_op_lerp_fused( GGML_ASSERT(grid <= std::numeric_limits::max()); + const int64_t base_total = ggml_nelements(x_prev); + const int64_t n_mix = dst->ne[3]; + + if (ggml_are_same_shape(x_prev, cur) && + dst->ne[0] == x_prev->ne[0] && + dst->ne[1] == x_prev->ne[1] && + dst->ne[2] == x_prev->ne[2] && + total == base_total * n_mix && + weight->ne[0] == dst->ne[0] && + weight->ne[1] == 1 && + weight->ne[2] == 1 && + weight->ne[3] == n_mix && + ggml_nelements(weight) == dst->ne[0] * n_mix) { + const int64_t base_grid = (base_total + block - 1) / block; + GGML_ASSERT(base_grid <= std::numeric_limits::max()); + GGML_ASSERT(n_mix <= std::numeric_limits::max()); + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) base_grid, (uint32_t) n_mix), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_lerp_fused_rwv_contig_f32, launch_params, + (const float *) x_prev->data, + (const float *) cur->data, + (const float *) weight->data, + (float *) dst->data, + dst->ne[0], + base_total); + return; + } + auto stride = [](const ggml_tensor * t, int dim) { return int64_t(t->nb[dim] / ggml_element_size(t)); }; From f5d7d00ad5c73cb582eff29a04f24935d01b31df Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 19:35:33 +0800 Subject: [PATCH 06/16] cuda: specialize rwkv7 wkv decode --- ggml/src/ggml-cuda/wkv.cu | 64 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index 3c104bfe2c23..c5b17d7827f1 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -154,6 +154,65 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons } } +template +static __global__ void __launch_bounds__(WARP_SIZE * rows_per_block, 2) +rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * r_k, const float * s, float * dst) { + constexpr int head_size = CUDA_WKV_BLOCK_SIZE; + constexpr int half_head = head_size / 2; + + const int lane = threadIdx.x; + const int row = blockIdx.y * rows_per_block + threadIdx.y; + const int bid = blockIdx.x; + + const int batch_i = bid / H; + const int head_i = bid % H; + const int state_size = C * head_size; + const int head_off = head_i * head_size; + const int t = batch_i * C + head_off + row; + + __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size], _r_k[head_size]; + + if (threadIdx.y == 0) { + constexpr float w_scale = -0.6065306597126334f; + + _r[lane] = r[batch_i * C + head_off + lane]; + _w[lane] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane]))); + _k[lane] = k[batch_i * C + head_off + lane]; + _kk[lane] = kk[batch_i * C + head_off + lane]; + _a[lane] = a[batch_i * C + head_off + lane]; + _r_k[lane] = r_k[head_off + lane]; + + _r[lane + half_head] = r[batch_i * C + head_off + lane + half_head]; + _w[lane + half_head] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane + half_head]))); + _k[lane + half_head] = k[batch_i * C + head_off + lane + half_head]; + _kk[lane + half_head] = kk[batch_i * C + head_off + lane + half_head]; + _a[lane + half_head] = a[batch_i * C + head_off + lane + half_head]; + _r_k[lane + half_head] = r_k[head_off + lane + half_head]; + } + __syncthreads(); + + const int64_t state_base = batch_i * state_size + head_i * head_size * head_size + row * head_size; + const float s0 = s[state_base + lane]; + const float s1 = s[state_base + lane + half_head]; + + const float sa = -warp_reduce_sum(_kk[lane] * s0 + _kk[lane + half_head] * s1); + const float rk = warp_reduce_sum(_r[lane] * _k[lane] * _r_k[lane] + + _r[lane + half_head] * _k[lane + half_head] * _r_k[lane + half_head]); + + const float vt = v[t]; + const float st0 = s0 * _w[lane] + _k[lane] * vt + sa * _kk[lane] * _a[lane]; + const float st1 = s1 * _w[lane + half_head] + _k[lane + half_head] * vt + sa * _kk[lane + half_head] * _a[lane + half_head]; + const float y = warp_reduce_sum(st0 * _r[lane] + st1 * _r[lane + half_head]); + + dst[2 * T * C + state_base + lane] = st0; + dst[2 * T * C + state_base + lane + half_head] = st1; + + if (lane == 0) { + dst[t] = y; + dst[T * C + t] = vt * rk; + } +} + void ggml_cuda_op_rwkv_wkv6(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float * k_d = (const float *)dst->src[0]->data; const float * v_d = (const float *)dst->src[1]->data; @@ -205,7 +264,10 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2); - if (C / H == CUDA_WKV_BLOCK_SIZE) { + if (T / B == 1 && C / H == CUDA_WKV_BLOCK_SIZE) { + constexpr int rows_per_block = 4; + rwkv_wkv7_f32_t1_warp_row<<>>(T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); + } else if (C / H == CUDA_WKV_BLOCK_SIZE) { rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); } else { rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); From 6a470d5eebd1c96e2315409ffacefdc9fb688f6e Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 21:19:23 +0800 Subject: [PATCH 07/16] vulkan: specialize rwkv7 decode and lerp --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 231 +++++++++++++++++- ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp | 33 +++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + .../ggml-vulkan/vulkan-shaders/wkv7_t1.comp | 103 ++++++++ 4 files changed, 368 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cc02f66fa58a..a2b1ebb93c81 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -814,6 +814,7 @@ struct vk_device_struct { vk_pipeline pipeline_mul[2][2][2]; vk_pipeline pipeline_mul_norepeat[2][2][2]; vk_pipeline pipeline_add_mul_f32; + vk_pipeline pipeline_lerp_f32; vk_pipeline pipeline_div[2][2][2]; vk_pipeline pipeline_div_norepeat[2][2][2]; vk_pipeline pipeline_add_rms[2][2][2]; @@ -939,6 +940,7 @@ struct vk_device_struct { vk_pipeline pipeline_pool2d_f32; vk_pipeline pipeline_rwkv_wkv6_f32; vk_pipeline pipeline_rwkv_wkv7_f32; + vk_pipeline pipeline_rwkv_wkv7_t1_f32; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; vk_pipeline pipeline_ssm_scan_f32_d128; @@ -1639,6 +1641,16 @@ struct vk_op_rwkv_wkv7_push_constants { uint32_t C; uint32_t H; }; + +struct vk_op_lerp_push_constants { + uint32_t ne0; + uint32_t base_total; + uint32_t x_offset; + uint32_t c_offset; + uint32_t w_offset; + uint32_t d_offset; +}; + struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -2153,6 +2165,7 @@ struct ggml_backend_vk_context { int fused_ops_write_mask {}; bool fused_add_mul {}; int fused_add_mul_dst_idx {-1}; + bool fused_lerp {}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; @@ -5079,6 +5092,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_add_mul_f32, "add_mul_f32", add_mul_f32_len, add_mul_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_lerp_f32, "lerp_f32", lerp_f32_len, lerp_f32_data, "main", 4, sizeof(vk_op_lerp_push_constants), {256, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true); for (auto &it : device->pipeline_fa_mask_opt) { @@ -5428,6 +5442,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 9, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); + if (device->subgroup_arithmetic && device->subgroup_size <= 64 && 64 % device->subgroup_size == 0) { + const uint32_t rows_per_wg = 4; + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_t1_f32, "rwkv_wkv7_t1_f32", rwkv_wkv7_t1_f32_len, rwkv_wkv7_t1_f32_data, "main", 9, sizeof(vk_op_rwkv_wkv7_push_constants), {1, rows_per_wg, 1}, {device->subgroup_size, rows_per_wg}, 1, true, true, device->subgroup_size); + } { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; @@ -11948,6 +11966,54 @@ static void ggml_vk_add_mul(ggml_backend_vk_context * ctx, vk_context& subctx, g ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, dst_buf, scale_buf }, pc, elements); } +static bool ggml_vk_should_fuse_lerp( + const ggml_backend_vk_context * ctx, + const ggml_tensor * sub, + const ggml_tensor * repeat, + const ggml_tensor * mul, + const ggml_tensor * add, + const ggml_tensor ** x_prev, + const ggml_tensor ** cur, + const ggml_tensor ** weight); + +static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { + const ggml_tensor * sub = cgraph->nodes[node_idx]; + const ggml_tensor * repeat = ctx->num_additional_fused_ops == 3 ? cgraph->nodes[node_idx + 1] : nullptr; + const ggml_tensor * mul = cgraph->nodes[node_idx + ctx->num_additional_fused_ops - 1]; + ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; + + const ggml_tensor * x_prev = nullptr; + const ggml_tensor * cur = nullptr; + const ggml_tensor * weight = nullptr; + GGML_ASSERT(ggml_vk_should_fuse_lerp(ctx, sub, repeat, mul, dst, &x_prev, &cur, &weight)); + + const uint32_t x_offset = get_misalign_bytes(ctx, x_prev) / ggml_type_size(x_prev->type); + const uint32_t c_offset = get_misalign_bytes(ctx, cur) / ggml_type_size(cur->type); + const uint32_t w_offset = get_misalign_bytes(ctx, weight) / ggml_type_size(weight->type); + const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + const uint32_t base_total = (uint32_t) ggml_nelements(x_prev); + const uint32_t n_mix = (uint32_t) (ggml_nelements(dst) / ggml_nelements(x_prev)); + + vk_op_lerp_push_constants pc { + (uint32_t) dst->ne[0], + base_total, + x_offset, + c_offset, + w_offset, + d_offset, + }; + + vk_subbuffer x_prev_buf = ggml_vk_tensor_subbuffer(ctx, x_prev, true); + vk_subbuffer cur_buf = ggml_vk_tensor_subbuffer(ctx, cur, true); + vk_subbuffer weight_buf = ggml_vk_tensor_subbuffer(ctx, weight, true); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_pipeline pipeline = ctx->device->pipeline_lerp_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_prev_buf, cur_buf, weight_buf, dst_buf }, pc, { base_total, n_mix, 1 }); +} + static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); @@ -12019,6 +12085,36 @@ static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx } } +static bool ggml_vk_rwkv_wkv7_t1(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants & pc) { + const uint32_t head_size = pc.H == 0 ? 0 : pc.C / pc.H; + if (pc.T != pc.B || head_size != 64 || ctx->device->pipeline_rwkv_wkv7_t1_f32 == nullptr) { + return false; + } + + for (int i = 0; i < 8; i++) { + GGML_ASSERT(!ggml_is_quantized(dst->src[i]->type)); + } + + GGML_ASSERT(dst->buffer != nullptr); + + vk_pipeline pipeline = ctx->device->pipeline_rwkv_wkv7_t1_f32; + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + vk_subbuffer src_buf[8] = {}; + for (int i = 0; i < 8; i++) { + src_buf[i] = ggml_vk_tensor_subbuffer(ctx, dst->src[i]); + } + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], src_buf[7], dst_buf}, + pc, {pc.B * pc.H, head_size, 1}); + + return true; +} + static void ggml_vk_rwkv_wkv6(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const size_t seq_length = dst->src[0]->ne[2]; const size_t n_embed = dst->ne[0]; @@ -12042,6 +12138,16 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx, const size_t n_embed = dst->ne[0]; const size_t n_heads = dst->src[0]->ne[1]; const size_t n_seqs = dst->src[7]->ne[1]; + const vk_op_rwkv_wkv6_push_constants pc = { + (uint32_t)n_seqs, + (uint32_t)seq_length, + (uint32_t)n_embed, + (uint32_t)n_heads, + }; + + if (ggml_vk_rwkv_wkv7_t1(ctx, subctx, dst, pc)) { + return; + } ggml_vk_op_f32_wkv( ctx, subctx, dst, @@ -14775,7 +14881,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr } break; case GGML_OP_SUB: - ggml_vk_sub(ctx, compute_ctx, src0, src1, node); + if (ctx->fused_lerp) { + ggml_vk_lerp_fused(ctx, compute_ctx, cgraph, node_idx); + } else { + ggml_vk_sub(ctx, compute_ctx, src0, src1, node); + } break; case GGML_OP_MUL: @@ -15861,6 +15971,100 @@ static bool ggml_vk_should_fuse_add_mul(const ggml_backend_vk_context * ctx, con return true; } +static bool ggml_vk_should_fuse_lerp( + const ggml_backend_vk_context * ctx, + const ggml_tensor * sub, + const ggml_tensor * repeat, + const ggml_tensor * mul, + const ggml_tensor * add, + const ggml_tensor ** x_prev, + const ggml_tensor ** cur, + const ggml_tensor ** weight) { + if (sub->op != GGML_OP_SUB || mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD) { + return false; + } + + const ggml_tensor * sx = sub; + if (repeat) { + if (repeat->op != GGML_OP_REPEAT || repeat->src[0] != sub) { + return false; + } + sx = repeat; + } + + const ggml_tensor * w = nullptr; + if (mul->src[0] == sx) { + w = mul->src[1]; + } else if (mul->src[1] == sx) { + w = mul->src[0]; + } else { + return false; + } + + const ggml_tensor * add_cur = nullptr; + if (add->src[0] == mul) { + add_cur = add->src[1]; + } else if (add->src[1] == mul) { + add_cur = add->src[0]; + } else { + return false; + } + + if (add_cur != sub->src[1]) { + return false; + } + + const ggml_tensor * xp = sub->src[0]; + const ggml_tensor * c = sub->src[1]; + if (xp->type != GGML_TYPE_F32 || c->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_can_repeat(xp, add) || !ggml_can_repeat(c, add) || !ggml_can_repeat(w, add)) { + return false; + } + + if (!ggml_is_contiguous(xp) || !ggml_is_contiguous(c) || !ggml_is_contiguous(w) || !ggml_is_contiguous(add)) { + return false; + } + + const int64_t base_total = ggml_nelements(xp); + const int64_t total = ggml_nelements(add); + if (!ggml_are_same_shape(xp, c) || + base_total <= 0 || + total % base_total != 0 || + add->ne[0] != xp->ne[0] || + add->ne[1] != xp->ne[1] || + add->ne[2] != xp->ne[2]) { + return false; + } + + const int64_t n_mix = total / base_total; + if (w->ne[0] != add->ne[0] || + w->ne[1] != 1 || + w->ne[2] != 1 || + w->ne[3] != n_mix || + ggml_nelements(w) != add->ne[0] * n_mix) { + return false; + } + + if (add->ne[0] > UINT32_MAX || base_total > UINT32_MAX || n_mix > UINT32_MAX) { + return false; + } + + if (get_misalign_bytes(ctx, xp) / ggml_type_size(xp->type) > UINT32_MAX || + get_misalign_bytes(ctx, c) / ggml_type_size(c->type) > UINT32_MAX || + get_misalign_bytes(ctx, w) / ggml_type_size(w->type) > UINT32_MAX || + get_misalign_bytes(ctx, add) / ggml_type_size(add->type) > UINT32_MAX) { + return false; + } + + *x_prev = xp; + *cur = c; + *weight = w; + return true; +} + static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; @@ -16556,6 +16760,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; + ctx->fused_lerp = false; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -16590,6 +16795,28 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; + } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_REPEAT, GGML_OP_MUL, GGML_OP_ADD }, { i + 3 })) { + const ggml_tensor * x_prev = nullptr; + const ggml_tensor * cur = nullptr; + const ggml_tensor * weight = nullptr; + if (ggml_vk_should_fuse_lerp(ctx, cgraph->nodes[i], cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], + &x_prev, &cur, &weight)) { + ctx->num_additional_fused_ops = 3; + ctx->fused_lerp = true; + fusion_string = "LERP"; + std::fill_n(op_srcs_fused_elementwise, 4, true); + } + } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_MUL, GGML_OP_ADD }, { i + 2 })) { + const ggml_tensor * x_prev = nullptr; + const ggml_tensor * cur = nullptr; + const ggml_tensor * weight = nullptr; + if (ggml_vk_should_fuse_lerp(ctx, cgraph->nodes[i], nullptr, cgraph->nodes[i + 1], cgraph->nodes[i + 2], + &x_prev, &cur, &weight)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_lerp = true; + fusion_string = "LERP"; + std::fill_n(op_srcs_fused_elementwise, 3, true); + } } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_ADD, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; ctx->fused_add_mul = true; @@ -16776,6 +17003,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; + ctx->fused_lerp = false; } } @@ -16830,6 +17058,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_ops_write_mask = 0; ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; + ctx->fused_lerp = false; } ctx->last_total_flops = total_flops; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp new file mode 100644 index 000000000000..fb52880527f0 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp @@ -0,0 +1,33 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint ne0; + uint base_total; + uint x_offset; + uint c_offset; + uint w_offset; + uint d_offset; +}; + +layout(binding = 0) readonly buffer XPrevBuf { A_TYPE x_prev[]; }; +layout(binding = 1) readonly buffer CurBuf { A_TYPE cur[]; }; +layout(binding = 2) readonly buffer WeightBuf { A_TYPE weight[]; }; +layout(binding = 3) buffer DstBuf { D_TYPE dst[]; }; + +void main() { + const uint ibase = gl_GlobalInvocationID.x; + if (ibase >= base_total) { + return; + } + + const uint imix = gl_GlobalInvocationID.y; + const uint i0 = ibase % ne0; + const A_TYPE c = cur[c_offset + ibase]; + + dst[d_offset + imix * base_total + ibase] = + D_TYPE(c + (x_prev[x_offset + ibase] - c) * weight[w_offset + imix * ne0 + i0]); +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index b2769b9ca41b..e5208e6b8131 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -837,6 +837,7 @@ void process_shaders() { string_to_spv("sub_f32", "sub.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + string_to_spv("lerp_f32", "lerp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); string_to_spv("add_mul_f32", "add_mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); @@ -1029,6 +1030,7 @@ void process_shaders() { string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + string_to_spv("rwkv_wkv7_t1_f32", "wkv7_t1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}})); string_to_spv("gated_delta_net_f32_nocluster", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "0"}})); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp new file mode 100644 index 000000000000..993bed63d8cf --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp @@ -0,0 +1,103 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 32; +layout(constant_id = 1) const uint ROWS_PER_WG = 4; + +const uint HEAD_SIZE = 64; +const uint ROWS_PER_LANE = HEAD_SIZE / SUBGROUP_SIZE; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint B; + uint T; + uint C; + uint H; +}; + +layout(binding = 0) readonly buffer RBuf { A_TYPE r[]; }; +layout(binding = 1) readonly buffer WBuf { A_TYPE w[]; }; +layout(binding = 2) readonly buffer KBuf { A_TYPE k[]; }; +layout(binding = 3) readonly buffer VBuf { A_TYPE v[]; }; +layout(binding = 4) readonly buffer KKBuf { A_TYPE kk[]; }; +layout(binding = 5) readonly buffer ABuf { A_TYPE a[]; }; +layout(binding = 6) readonly buffer RKBuf { A_TYPE r_k[]; }; +layout(binding = 7) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 8) buffer DstBuf { A_TYPE dst[]; }; + +shared A_TYPE _r[HEAD_SIZE], _w[HEAD_SIZE], _k[HEAD_SIZE], _kk[HEAD_SIZE], _a[HEAD_SIZE], _r_k[HEAD_SIZE]; + +void main() { + const uint batch_id = gl_WorkGroupID.x / H; + const uint head_id = gl_WorkGroupID.x % H; + const uint row = gl_WorkGroupID.y * ROWS_PER_WG + gl_LocalInvocationID.y; + const uint lane = gl_SubgroupInvocationID; + + if (batch_id >= B || head_id >= H || row >= HEAD_SIZE) { + return; + } + + const uint state_size = C * HEAD_SIZE; + const uint head_off = head_id * HEAD_SIZE; + const uint token_base = batch_id * C + head_off; + + if (gl_LocalInvocationID.y == 0) { + [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { + const uint i = r_i * SUBGROUP_SIZE + lane; + const uint t = token_base + i; + const A_TYPE w_scale = A_TYPE(-0.6065306597126334); + _r[i] = r[t]; + _w[i] = exp(w_scale / (A_TYPE(1.0) + exp(-w[t]))); + _k[i] = k[t]; + _kk[i] = kk[t]; + _a[i] = a[t]; + _r_k[i] = r_k[head_off + i]; + } + } + barrier(); + + A_TYPE s_shard[ROWS_PER_LANE]; + [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { + const uint i = r_i * SUBGROUP_SIZE + lane; + s_shard[r_i] = state_in[batch_id * state_size + head_id * HEAD_SIZE * HEAD_SIZE + row * HEAD_SIZE + i]; + } + + A_TYPE sa_partial = 0.0; + A_TYPE rk_partial = 0.0; + [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { + const uint i = r_i * SUBGROUP_SIZE + lane; + sa_partial += s_shard[r_i] * _kk[i]; + rk_partial += _r[i] * _k[i] * _r_k[i]; + } + + const A_TYPE sa = -subgroupAdd(sa_partial); + const A_TYPE rk = subgroupAdd(rk_partial); + + const uint t_out = token_base + row; + const A_TYPE v_val = v[t_out]; + A_TYPE y_partial = 0.0; + + [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { + const uint i = r_i * SUBGROUP_SIZE + lane; + A_TYPE s = s_shard[r_i]; + s = s * _w[i] + _k[i] * v_val + sa * _kk[i] * _a[i]; + s_shard[r_i] = s; + y_partial += _r[i] * s; + } + + const A_TYPE y = subgroupAdd(y_partial); + + if (lane == 0) { + dst[t_out] = y; + dst[T * C + t_out] = v_val * rk; + } + + [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { + const uint i = r_i * SUBGROUP_SIZE + lane; + dst[2 * T * C + batch_id * state_size + head_id * HEAD_SIZE * HEAD_SIZE + row * HEAD_SIZE + i] = s_shard[r_i]; + } +} From 402d7a70ef22fd663bfe555cce4851f14b6df88d Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Wed, 1 Jul 2026 22:34:08 +0800 Subject: [PATCH 08/16] rwkv: fuse rk correction separately --- ggml/include/ggml.h | 1 - ggml/src/ggml-cpu/ops.cpp | 39 +- ggml/src/ggml-cuda/binbcast.cu | 290 ------------- ggml/src/ggml-cuda/binbcast.cuh | 21 - ggml/src/ggml-cuda/fused-ops.cu | 385 ++++++++++++++++++ ggml/src/ggml-cuda/fused-ops.cuh | 31 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 155 +++++++ ggml/src/ggml-cuda/wkv.cu | 42 +- ggml/src/ggml-metal/ggml-metal-device.cpp | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 3 +- ggml/src/ggml-metal/ggml-metal.metal | 14 +- ggml/src/ggml-sycl/wkv.cpp | 27 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 337 ++++++++++++++- .../ggml-vulkan/vulkan-shaders/rwkv_rk.comp | 49 +++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp | 16 +- .../ggml-vulkan/vulkan-shaders/wkv7_t1.comp | 14 +- ggml/src/ggml.c | 12 +- src/models/rwkv7-base.cpp | 8 +- tests/test-backend-ops.cpp | 3 +- 20 files changed, 1002 insertions(+), 448 deletions(-) create mode 100644 ggml/src/ggml-cuda/fused-ops.cu create mode 100644 ggml/src/ggml-cuda/fused-ops.cuh create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/rwkv_rk.comp diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index f3839e1a5688..4f96588404ae 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2527,7 +2527,6 @@ extern "C" { struct ggml_tensor * v, struct ggml_tensor * kk, struct ggml_tensor * a, - struct ggml_tensor * r_k, struct ggml_tensor * state); /* Solves a specific equation of the form Ax=B, where A is a triangular matrix diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 4e41c2a1f93a..bc60564ab281 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11,7 +11,6 @@ #include #include #include -#include // ggml_compute_forward_dup @@ -10962,22 +10961,17 @@ void ggml_compute_forward_gated_delta_net( // ggml_compute_forward_rwkv_wkv7 -static inline float ggml_rwkv_wkv7_decay(float w) { - constexpr float w_scale = -0.6065306597126334f; - return expf(w_scale / (1.0f + expf(-w))); -} - static void ggml_compute_forward_rwkv_wkv7_f32( const ggml_compute_params * params, ggml_tensor * dst) { const int64_t T = dst->src[1]->ne[2]; const int64_t C = dst->ne[0]; const int64_t HEADS = dst->src[1]->ne[1]; - const int64_t n_seqs = dst->src[7]->ne[1]; + const int64_t n_seqs = dst->src[6]->ne[1]; const int64_t head_size = C / HEADS; float * dst_data = (float *) dst->data; - float * state = ((float *) dst->data) + 2 * C * T; + float * state = ((float *) dst->data) + C * T; const int ith = params->ith; const int nth = params->nth; @@ -10992,14 +10986,14 @@ static void ggml_compute_forward_rwkv_wkv7_f32( float * v = (float *) dst->src[3]->data; float * kk = (float *) dst->src[4]->data; float * a = (float *) dst->src[5]->data; - float * r_k = (float *) dst->src[6]->data; int64_t t_stride = HEADS * head_size; // Same to C int64_t h_stride = C / HEADS; GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS + GGML_ASSERT(head_size == 64 || head_size == 128); int64_t h_stride_2d = head_size * head_size; - std::vector w_decay(head_size); + float w_decay[128]; #if defined(GGML_SIMD) #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v_intrinsic) @@ -11008,18 +11002,17 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; - float rk = 0; + constexpr float w_scale = -0.6065306597126334f; for (int64_t j = 0; j < head_size; j++) { const int64_t t_h_j_offset = t_h_offset + j; - rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; - w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + w_decay[j] = expf(w_scale / (1.0f + expf(-w[t_h_j_offset]))); } for (int64_t i = 0; i < head_size; i++) { @@ -11048,7 +11041,6 @@ static void ggml_compute_forward_rwkv_wkv7_f32( result += state_cur[h_2d_i_j_offset] * r_val; } dst_data[t_h_i_offset] = result; - dst_data[C * T + t_h_i_offset] = v_val * rk; } } } @@ -11057,18 +11049,17 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; - float rk = 0; + constexpr float w_scale = -0.6065306597126334f; for (int64_t j = 0; j < head_size; j++) { const int64_t t_h_j_offset = t_h_offset + j; - rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; - w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + w_decay[j] = expf(w_scale / (1.0f + expf(-w[t_h_j_offset]))); } for (int64_t ii = 0; ii < head_size; ii++) { @@ -11121,7 +11112,6 @@ static void ggml_compute_forward_rwkv_wkv7_f32( } } GGML_F32_VEC_REDUCE(dst_data[t_h_i_offset], result_vec); - dst_data[C * T + t_h_i_offset] = v[t_h_i_offset] * rk; // There shouldn't be left-overs though. for (; j < head_size; j++) { @@ -11137,7 +11127,6 @@ static void ggml_compute_forward_rwkv_wkv7_f32( float prev_state_val = state_prev[h_2d_i_j_offset]; state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; dst_data[t_h_i_offset] += state_cur[h_2d_i_j_offset] * r_val; - dst_data[C * T + t_h_i_offset] = v[t_h_i_offset] * rk; } } } @@ -11148,18 +11137,17 @@ static void ggml_compute_forward_rwkv_wkv7_f32( int64_t t_offset = t * t_stride; int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[7]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { int64_t h_offset = h * h_stride; int64_t t_h_offset = t_offset + h_offset; int64_t h_2d_offset = h * h_stride_2d; - float rk = 0; + constexpr float w_scale = -0.6065306597126334f; for (int64_t j = 0; j < head_size; j++) { const int64_t t_h_j_offset = t_h_offset + j; - rk += k[t_h_j_offset] * r[t_h_j_offset] * r_k[h_offset + j]; - w_decay[j] = ggml_rwkv_wkv7_decay(w[t_h_j_offset]); + w_decay[j] = expf(w_scale / (1.0f + expf(-w[t_h_j_offset]))); } for (int64_t i = 0; i < head_size; i++) { @@ -11188,7 +11176,6 @@ static void ggml_compute_forward_rwkv_wkv7_f32( result += state_cur[h_2d_i_j_offset] * r_val; } dst_data[t_h_i_offset] = result; - dst_data[C * T + t_h_i_offset] = v_val * rk; } } } diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index fa0d03ec55e7..839c32a48aa7 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -26,137 +26,6 @@ static __device__ __forceinline__ float op_div(const float a, const float b) { return a / b; } -static __device__ __forceinline__ size_t ggml_cuda_index_4d( - const int64_t i0, const int64_t i1, const int64_t i2, const int64_t i3, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, - const int64_t s0, const int64_t s1, const int64_t s2, const int64_t s3) { - const int64_t j0 = i0 % ne0; - const int64_t j1 = i1 % ne1; - const int64_t j2 = i2 % ne2; - const int64_t j3 = i3 % ne3; - - return size_t(j0)*s0 + size_t(j1)*s1 + size_t(j2)*s2 + size_t(j3)*s3; -} - -static __global__ void k_lerp_fused_f32( - const float * __restrict__ x_prev, - const float * __restrict__ cur, - const float * __restrict__ weight, - float * __restrict__ dst, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, - const int64_t xp_ne0, const int64_t xp_ne1, const int64_t xp_ne2, const int64_t xp_ne3, - const int64_t xp_s0, const int64_t xp_s1, const int64_t xp_s2, const int64_t xp_s3, - const int64_t c_ne0, const int64_t c_ne1, const int64_t c_ne2, const int64_t c_ne3, - const int64_t c_s0, const int64_t c_s1, const int64_t c_s2, const int64_t c_s3, - const int64_t w_ne0, const int64_t w_ne1, const int64_t w_ne2, const int64_t w_ne3, - const int64_t w_s0, const int64_t w_s1, const int64_t w_s2, const int64_t w_s3) { - const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = ne0 * ne1 * ne2 * ne3; - - if (idx >= total) { - return; - } - - const int64_t i0 = idx % ne0; - const int64_t t1 = idx / ne0; - const int64_t i1 = t1 % ne1; - const int64_t t2 = t1 / ne1; - const int64_t i2 = t2 % ne2; - const int64_t i3 = t2 / ne2; - - const size_t ixp = ggml_cuda_index_4d(i0, i1, i2, i3, xp_ne0, xp_ne1, xp_ne2, xp_ne3, xp_s0, xp_s1, xp_s2, xp_s3); - const size_t ic = ggml_cuda_index_4d(i0, i1, i2, i3, c_ne0, c_ne1, c_ne2, c_ne3, c_s0, c_s1, c_s2, c_s3); - const size_t iw = ggml_cuda_index_4d(i0, i1, i2, i3, w_ne0, w_ne1, w_ne2, w_ne3, w_s0, w_s1, w_s2, w_s3); - - const float c = cur[ic]; - dst[idx] = c + (x_prev[ixp] - c) * weight[iw]; -} - -static __global__ void k_lerp_fused_rwv_contig_f32( - const float * __restrict__ x_prev, - const float * __restrict__ cur, - const float * __restrict__ weight, - float * __restrict__ dst, - const int64_t ne0, - const int64_t base_total) { - const int64_t ibase = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - if (ibase >= base_total) { - return; - } - - const int64_t imix = blockIdx.y; - const int64_t i0 = ibase % ne0; - const float c = cur[ibase]; - dst[imix * base_total + ibase] = c + (x_prev[ibase] - c) * weight[imix * ne0 + i0]; -} - -static __global__ void k_mul_sub_add_fused_f32( - const float * __restrict__ base, - const float * __restrict__ scale, - const float * __restrict__ value, - float * __restrict__ dst, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, - const int64_t b_ne0, const int64_t b_ne1, const int64_t b_ne2, const int64_t b_ne3, - const int64_t b_s0, const int64_t b_s1, const int64_t b_s2, const int64_t b_s3, - const int64_t s_ne0, const int64_t s_ne1, const int64_t s_ne2, const int64_t s_ne3, - const int64_t s_s0, const int64_t s_s1, const int64_t s_s2, const int64_t s_s3, - const int64_t v_ne0, const int64_t v_ne1, const int64_t v_ne2, const int64_t v_ne3, - const int64_t v_s0, const int64_t v_s1, const int64_t v_s2, const int64_t v_s3) { - const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = ne0 * ne1 * ne2 * ne3; - - if (idx >= total) { - return; - } - - const int64_t i0 = idx % ne0; - const int64_t t1 = idx / ne0; - const int64_t i1 = t1 % ne1; - const int64_t t2 = t1 / ne1; - const int64_t i2 = t2 % ne2; - const int64_t i3 = t2 / ne2; - - const size_t ib = ggml_cuda_index_4d(i0, i1, i2, i3, b_ne0, b_ne1, b_ne2, b_ne3, b_s0, b_s1, b_s2, b_s3); - const size_t is = ggml_cuda_index_4d(i0, i1, i2, i3, s_ne0, s_ne1, s_ne2, s_ne3, s_s0, s_s1, s_s2, s_s3); - const size_t iv = ggml_cuda_index_4d(i0, i1, i2, i3, v_ne0, v_ne1, v_ne2, v_ne3, v_s0, v_s1, v_s2, v_s3); - - const float v = value[iv]; - dst[idx] = base[ib] + (scale[is] * v - v); -} - -static __global__ void k_add_mul_fused_f32( - const float * __restrict__ src0, - const float * __restrict__ src1, - const float * __restrict__ scale, - float * __restrict__ dst, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, - const int64_t s0_ne0, const int64_t s0_ne1, const int64_t s0_ne2, const int64_t s0_ne3, - const int64_t s0_s0, const int64_t s0_s1, const int64_t s0_s2, const int64_t s0_s3, - const int64_t s1_ne0, const int64_t s1_ne1, const int64_t s1_ne2, const int64_t s1_ne3, - const int64_t s1_s0, const int64_t s1_s1, const int64_t s1_s2, const int64_t s1_s3, - const int64_t sc_ne0, const int64_t sc_ne1, const int64_t sc_ne2, const int64_t sc_ne3, - const int64_t sc_s0, const int64_t sc_s1, const int64_t sc_s2, const int64_t sc_s3) { - const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = ne0 * ne1 * ne2 * ne3; - - if (idx >= total) { - return; - } - - const int64_t i0 = idx % ne0; - const int64_t t1 = idx / ne0; - const int64_t i1 = t1 % ne1; - const int64_t t2 = t1 / ne1; - const int64_t i2 = t2 % ne2; - const int64_t i3 = t2 / ne2; - - const size_t is0 = ggml_cuda_index_4d(i0, i1, i2, i3, s0_ne0, s0_ne1, s0_ne2, s0_ne3, s0_s0, s0_s1, s0_s2, s0_s3); - const size_t is1 = ggml_cuda_index_4d(i0, i1, i2, i3, s1_ne0, s1_ne1, s1_ne2, s1_ne3, s1_s0, s1_s1, s1_s2, s1_s3); - const size_t isc = ggml_cuda_index_4d(i0, i1, i2, i3, sc_ne0, sc_ne1, sc_ne2, sc_ne3, sc_s0, sc_s1, sc_s2, sc_s3); - - dst[idx] = (src0[is0] + src1[is1]) * scale[isc]; -} - template type == GGML_TYPE_F32); - GGML_ASSERT(cur->type == GGML_TYPE_F32); - GGML_ASSERT(weight->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_can_repeat(x_prev, dst)); - GGML_ASSERT(ggml_can_repeat(cur, dst)); - GGML_ASSERT(ggml_can_repeat(weight, dst)); - GGML_ASSERT(ggml_is_contiguous(x_prev)); - GGML_ASSERT(ggml_is_contiguous(cur)); - GGML_ASSERT(ggml_is_contiguous(weight)); - GGML_ASSERT(ggml_is_contiguous(dst)); - - const int64_t total = ggml_nelements(dst); - const int block = 256; - const int64_t grid = (total + block - 1) / block; - - GGML_ASSERT(grid <= std::numeric_limits::max()); - - const int64_t base_total = ggml_nelements(x_prev); - const int64_t n_mix = dst->ne[3]; - - if (ggml_are_same_shape(x_prev, cur) && - dst->ne[0] == x_prev->ne[0] && - dst->ne[1] == x_prev->ne[1] && - dst->ne[2] == x_prev->ne[2] && - total == base_total * n_mix && - weight->ne[0] == dst->ne[0] && - weight->ne[1] == 1 && - weight->ne[2] == 1 && - weight->ne[3] == n_mix && - ggml_nelements(weight) == dst->ne[0] * n_mix) { - const int64_t base_grid = (base_total + block - 1) / block; - GGML_ASSERT(base_grid <= std::numeric_limits::max()); - GGML_ASSERT(n_mix <= std::numeric_limits::max()); - - const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) base_grid, (uint32_t) n_mix), block, 0, ctx.stream()); - ggml_cuda_kernel_launch(k_lerp_fused_rwv_contig_f32, launch_params, - (const float *) x_prev->data, - (const float *) cur->data, - (const float *) weight->data, - (float *) dst->data, - dst->ne[0], - base_total); - return; - } - - auto stride = [](const ggml_tensor * t, int dim) { - return int64_t(t->nb[dim] / ggml_element_size(t)); - }; - - const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); - ggml_cuda_kernel_launch(k_lerp_fused_f32, launch_params, - (const float *) x_prev->data, - (const float *) cur->data, - (const float *) weight->data, - (float *) dst->data, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - x_prev->ne[0], x_prev->ne[1], x_prev->ne[2], x_prev->ne[3], - stride(x_prev, 0), stride(x_prev, 1), stride(x_prev, 2), stride(x_prev, 3), - cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3], - stride(cur, 0), stride(cur, 1), stride(cur, 2), stride(cur, 3), - weight->ne[0], weight->ne[1], weight->ne[2], weight->ne[3], - stride(weight, 0), stride(weight, 1), stride(weight, 2), stride(weight, 3)); -} - -void ggml_cuda_op_mul_sub_add_fused( - ggml_backend_cuda_context & ctx, - const ggml_tensor * base, - const ggml_tensor * scale, - const ggml_tensor * value, - ggml_tensor * dst) { - GGML_ASSERT(base->type == GGML_TYPE_F32); - GGML_ASSERT(scale->type == GGML_TYPE_F32); - GGML_ASSERT(value->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_can_repeat(base, dst)); - GGML_ASSERT(ggml_can_repeat(scale, dst)); - GGML_ASSERT(ggml_can_repeat(value, dst)); - GGML_ASSERT(ggml_is_contiguous(base)); - GGML_ASSERT(ggml_is_contiguous(scale)); - GGML_ASSERT(ggml_is_contiguous(value)); - GGML_ASSERT(ggml_is_contiguous(dst)); - - const int64_t total = ggml_nelements(dst); - const int block = 256; - const int64_t grid = (total + block - 1) / block; - - GGML_ASSERT(grid <= std::numeric_limits::max()); - - auto stride = [](const ggml_tensor * t, int dim) { - return int64_t(t->nb[dim] / ggml_element_size(t)); - }; - - const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); - ggml_cuda_kernel_launch(k_mul_sub_add_fused_f32, launch_params, - (const float *) base->data, - (const float *) scale->data, - (const float *) value->data, - (float *) dst->data, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - base->ne[0], base->ne[1], base->ne[2], base->ne[3], - stride(base, 0), stride(base, 1), stride(base, 2), stride(base, 3), - scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], - stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3), - value->ne[0], value->ne[1], value->ne[2], value->ne[3], - stride(value, 0), stride(value, 1), stride(value, 2), stride(value, 3)); -} - -void ggml_cuda_op_add_mul_fused( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, - const ggml_tensor * src1, - const ggml_tensor * scale, - ggml_tensor * dst) { - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(scale->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_can_repeat(src0, dst)); - GGML_ASSERT(ggml_can_repeat(src1, dst)); - GGML_ASSERT(ggml_can_repeat(scale, dst)); - GGML_ASSERT(ggml_is_contiguous(src0)); - GGML_ASSERT(ggml_is_contiguous(src1)); - GGML_ASSERT(ggml_is_contiguous(scale)); - GGML_ASSERT(ggml_is_contiguous(dst)); - - const int64_t total = ggml_nelements(dst); - const int block = 256; - const int64_t grid = (total + block - 1) / block; - - GGML_ASSERT(grid <= std::numeric_limits::max()); - - auto stride = [](const ggml_tensor * t, int dim) { - return int64_t(t->nb[dim] / ggml_element_size(t)); - }; - - const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); - ggml_cuda_kernel_launch(k_add_mul_fused_f32, launch_params, - (const float *) src0->data, - (const float *) src1->data, - (const float *) scale->data, - (float *) dst->data, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], - stride(src0, 0), stride(src0, 1), stride(src0, 2), stride(src0, 3), - src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - stride(src1, 0), stride(src1, 1), stride(src1, 2), stride(src1, 3), - scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], - stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3)); -} void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 40ec517f2e1f..12624785b444 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -10,24 +10,3 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); - -void ggml_cuda_op_lerp_fused( - ggml_backend_cuda_context & ctx, - const ggml_tensor * x_prev, - const ggml_tensor * cur, - const ggml_tensor * weight, - ggml_tensor * dst); - -void ggml_cuda_op_mul_sub_add_fused( - ggml_backend_cuda_context & ctx, - const ggml_tensor * base, - const ggml_tensor * scale, - const ggml_tensor * value, - ggml_tensor * dst); - -void ggml_cuda_op_add_mul_fused( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, - const ggml_tensor * src1, - const ggml_tensor * scale, - ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/fused-ops.cu b/ggml/src/ggml-cuda/fused-ops.cu new file mode 100644 index 000000000000..74ae905eda45 --- /dev/null +++ b/ggml/src/ggml-cuda/fused-ops.cu @@ -0,0 +1,385 @@ +#include "fused-ops.cuh" + +#include +#include + +static __device__ __forceinline__ size_t ggml_cuda_index_4d( + const int64_t i0, const int64_t i1, const int64_t i2, const int64_t i3, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s0, const int64_t s1, const int64_t s2, const int64_t s3) { + const int64_t j0 = i0 % ne0; + const int64_t j1 = i1 % ne1; + const int64_t j2 = i2 % ne2; + const int64_t j3 = i3 % ne3; + + return size_t(j0)*s0 + size_t(j1)*s1 + size_t(j2)*s2 + size_t(j3)*s3; +} + +static __global__ void k_lerp_fused_f32( + const float * __restrict__ x_prev, + const float * __restrict__ cur, + const float * __restrict__ weight, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t xp_ne0, const int64_t xp_ne1, const int64_t xp_ne2, const int64_t xp_ne3, + const int64_t xp_s0, const int64_t xp_s1, const int64_t xp_s2, const int64_t xp_s3, + const int64_t c_ne0, const int64_t c_ne1, const int64_t c_ne2, const int64_t c_ne3, + const int64_t c_s0, const int64_t c_s1, const int64_t c_s2, const int64_t c_s3, + const int64_t w_ne0, const int64_t w_ne1, const int64_t w_ne2, const int64_t w_ne3, + const int64_t w_s0, const int64_t w_s1, const int64_t w_s2, const int64_t w_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t ixp = ggml_cuda_index_4d(i0, i1, i2, i3, xp_ne0, xp_ne1, xp_ne2, xp_ne3, xp_s0, xp_s1, xp_s2, xp_s3); + const size_t ic = ggml_cuda_index_4d(i0, i1, i2, i3, c_ne0, c_ne1, c_ne2, c_ne3, c_s0, c_s1, c_s2, c_s3); + const size_t iw = ggml_cuda_index_4d(i0, i1, i2, i3, w_ne0, w_ne1, w_ne2, w_ne3, w_s0, w_s1, w_s2, w_s3); + + const float c = cur[ic]; + dst[idx] = c + (x_prev[ixp] - c) * weight[iw]; +} + +static __global__ void k_lerp_fused_rwv_contig_f32( + const float * __restrict__ x_prev, + const float * __restrict__ cur, + const float * __restrict__ weight, + float * __restrict__ dst, + const int64_t ne0, + const int64_t base_total) { + const int64_t ibase = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (ibase >= base_total) { + return; + } + + const int64_t imix = blockIdx.y; + const int64_t i0 = ibase % ne0; + const float c = cur[ibase]; + dst[imix * base_total + ibase] = c + (x_prev[ibase] - c) * weight[imix * ne0 + i0]; +} + +static __global__ void k_mul_sub_add_fused_f32( + const float * __restrict__ base, + const float * __restrict__ scale, + const float * __restrict__ value, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t b_ne0, const int64_t b_ne1, const int64_t b_ne2, const int64_t b_ne3, + const int64_t b_s0, const int64_t b_s1, const int64_t b_s2, const int64_t b_s3, + const int64_t s_ne0, const int64_t s_ne1, const int64_t s_ne2, const int64_t s_ne3, + const int64_t s_s0, const int64_t s_s1, const int64_t s_s2, const int64_t s_s3, + const int64_t v_ne0, const int64_t v_ne1, const int64_t v_ne2, const int64_t v_ne3, + const int64_t v_s0, const int64_t v_s1, const int64_t v_s2, const int64_t v_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t ib = ggml_cuda_index_4d(i0, i1, i2, i3, b_ne0, b_ne1, b_ne2, b_ne3, b_s0, b_s1, b_s2, b_s3); + const size_t is = ggml_cuda_index_4d(i0, i1, i2, i3, s_ne0, s_ne1, s_ne2, s_ne3, s_s0, s_s1, s_s2, s_s3); + const size_t iv = ggml_cuda_index_4d(i0, i1, i2, i3, v_ne0, v_ne1, v_ne2, v_ne3, v_s0, v_s1, v_s2, v_s3); + + const float v = value[iv]; + dst[idx] = base[ib] + (scale[is] * v - v); +} + +static __global__ void k_add_mul_fused_f32( + const float * __restrict__ src0, + const float * __restrict__ src1, + const float * __restrict__ scale, + float * __restrict__ dst, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t s0_ne0, const int64_t s0_ne1, const int64_t s0_ne2, const int64_t s0_ne3, + const int64_t s0_s0, const int64_t s0_s1, const int64_t s0_s2, const int64_t s0_s3, + const int64_t s1_ne0, const int64_t s1_ne1, const int64_t s1_ne2, const int64_t s1_ne3, + const int64_t s1_s0, const int64_t s1_s1, const int64_t s1_s2, const int64_t s1_s3, + const int64_t sc_ne0, const int64_t sc_ne1, const int64_t sc_ne2, const int64_t sc_ne3, + const int64_t sc_s0, const int64_t sc_s1, const int64_t sc_s2, const int64_t sc_s3) { + const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = ne0 * ne1 * ne2 * ne3; + + if (idx >= total) { + return; + } + + const int64_t i0 = idx % ne0; + const int64_t t1 = idx / ne0; + const int64_t i1 = t1 % ne1; + const int64_t t2 = t1 / ne1; + const int64_t i2 = t2 % ne2; + const int64_t i3 = t2 / ne2; + + const size_t is0 = ggml_cuda_index_4d(i0, i1, i2, i3, s0_ne0, s0_ne1, s0_ne2, s0_ne3, s0_s0, s0_s1, s0_s2, s0_s3); + const size_t is1 = ggml_cuda_index_4d(i0, i1, i2, i3, s1_ne0, s1_ne1, s1_ne2, s1_ne3, s1_s0, s1_s1, s1_s2, s1_s3); + const size_t isc = ggml_cuda_index_4d(i0, i1, i2, i3, sc_ne0, sc_ne1, sc_ne2, sc_ne3, sc_s0, sc_s1, sc_s2, sc_s3); + + dst[idx] = (src0[is0] + src1[is1]) * scale[isc]; +} + +template +static __global__ void k_rwkv_rk_fused_f32( + const float * __restrict__ cur, + const float * __restrict__ k, + const float * __restrict__ r, + const float * __restrict__ v, + const float * __restrict__ r_k, + float * __restrict__ dst, + const int64_t C, + const int64_t H) { + const int64_t tid = threadIdx.x; + const int64_t h = blockIdx.x % H; + const int64_t t = blockIdx.x / H; + const int64_t off = t*C + h*head_size; + + __shared__ float rk[head_size]; + rk[tid] = k[off + tid] * r[off + tid] * r_k[h*head_size + tid]; + __syncthreads(); + + for (int stride = head_size / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + rk[tid] += rk[tid + stride]; + } + __syncthreads(); + } + + dst[off + tid] = cur[off + tid] + v[off + tid] * rk[0]; +} + +void ggml_cuda_op_lerp_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * x_prev, + const ggml_tensor * cur, + const ggml_tensor * weight, + ggml_tensor * dst) { + GGML_ASSERT(x_prev->type == GGML_TYPE_F32); + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(weight->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(x_prev, dst)); + GGML_ASSERT(ggml_can_repeat(cur, dst)); + GGML_ASSERT(ggml_can_repeat(weight, dst)); + GGML_ASSERT(ggml_is_contiguous(x_prev)); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(weight)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + const int64_t base_total = ggml_nelements(x_prev); + const int64_t n_mix = dst->ne[3]; + + if (ggml_are_same_shape(x_prev, cur) && + dst->ne[0] == x_prev->ne[0] && + dst->ne[1] == x_prev->ne[1] && + dst->ne[2] == x_prev->ne[2] && + total == base_total * n_mix && + weight->ne[0] == dst->ne[0] && + weight->ne[1] == 1 && + weight->ne[2] == 1 && + weight->ne[3] == n_mix && + ggml_nelements(weight) == dst->ne[0] * n_mix) { + const int64_t base_grid = (base_total + block - 1) / block; + GGML_ASSERT(base_grid <= std::numeric_limits::max()); + GGML_ASSERT(n_mix <= std::numeric_limits::max()); + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) base_grid, (uint32_t) n_mix), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_lerp_fused_rwv_contig_f32, launch_params, + (const float *) x_prev->data, + (const float *) cur->data, + (const float *) weight->data, + (float *) dst->data, + dst->ne[0], + base_total); + return; + } + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_lerp_fused_f32, launch_params, + (const float *) x_prev->data, + (const float *) cur->data, + (const float *) weight->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + x_prev->ne[0], x_prev->ne[1], x_prev->ne[2], x_prev->ne[3], + stride(x_prev, 0), stride(x_prev, 1), stride(x_prev, 2), stride(x_prev, 3), + cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3], + stride(cur, 0), stride(cur, 1), stride(cur, 2), stride(cur, 3), + weight->ne[0], weight->ne[1], weight->ne[2], weight->ne[3], + stride(weight, 0), stride(weight, 1), stride(weight, 2), stride(weight, 3)); +} + +void ggml_cuda_op_mul_sub_add_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * base, + const ggml_tensor * scale, + const ggml_tensor * value, + ggml_tensor * dst) { + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(value->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(base, dst)); + GGML_ASSERT(ggml_can_repeat(scale, dst)); + GGML_ASSERT(ggml_can_repeat(value, dst)); + GGML_ASSERT(ggml_is_contiguous(base)); + GGML_ASSERT(ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(value)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_mul_sub_add_fused_f32, launch_params, + (const float *) base->data, + (const float *) scale->data, + (const float *) value->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + base->ne[0], base->ne[1], base->ne[2], base->ne[3], + stride(base, 0), stride(base, 1), stride(base, 2), stride(base, 3), + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], + stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3), + value->ne[0], value->ne[1], value->ne[2], value->ne[3], + stride(value, 0), stride(value, 1), stride(value, 2), stride(value, 3)); +} + +void ggml_cuda_op_add_mul_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * scale, + ggml_tensor * dst) { + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_can_repeat(src0, dst)); + GGML_ASSERT(ggml_can_repeat(src1, dst)); + GGML_ASSERT(ggml_can_repeat(scale, dst)); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(scale)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t total = ggml_nelements(dst); + const int block = 256; + const int64_t grid = (total + block - 1) / block; + + GGML_ASSERT(grid <= std::numeric_limits::max()); + + auto stride = [](const ggml_tensor * t, int dim) { + return int64_t(t->nb[dim] / ggml_element_size(t)); + }; + + const ggml_cuda_kernel_launch_params launch_params(dim3((uint32_t) grid), block, 0, ctx.stream()); + ggml_cuda_kernel_launch(k_add_mul_fused_f32, launch_params, + (const float *) src0->data, + (const float *) src1->data, + (const float *) scale->data, + (float *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], + stride(src0, 0), stride(src0, 1), stride(src0, 2), stride(src0, 3), + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + stride(src1, 0), stride(src1, 1), stride(src1, 2), stride(src1, 3), + scale->ne[0], scale->ne[1], scale->ne[2], scale->ne[3], + stride(scale, 0), stride(scale, 1), stride(scale, 2), stride(scale, 3)); +} + +void ggml_cuda_op_rwkv_rk_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * cur, + const ggml_tensor * k, + const ggml_tensor * r, + const ggml_tensor * v, + const ggml_tensor * r_k, + ggml_tensor * dst) { + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(r->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(r_k->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r_k)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t head_size = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t T = k->ne[2]; + const int64_t C = head_size * H; + + GGML_ASSERT(head_size == 64 || head_size == 128); + GGML_ASSERT(dst->ne[0] == C && dst->ne[1] == T); + GGML_ASSERT(cur->ne[0] == C && cur->ne[1] == T); + GGML_ASSERT(r_k->ne[0] == head_size && r_k->ne[1] == H); + GGML_ASSERT(ggml_are_same_shape(k, r)); + GGML_ASSERT(ggml_are_same_shape(k, v)); + + const int64_t grid = H * T; + GGML_ASSERT(grid <= std::numeric_limits::max()); + + const cudaStream_t stream = ctx.stream(); + if (head_size == 64) { + const ggml_cuda_kernel_launch_params launch_params((uint32_t) grid, 64, 0, stream); + ggml_cuda_kernel_launch(k_rwkv_rk_fused_f32<64>, launch_params, + (const float *) cur->data, + (const float *) k->data, + (const float *) r->data, + (const float *) v->data, + (const float *) r_k->data, + (float *) dst->data, + C, H); + } else { + const ggml_cuda_kernel_launch_params launch_params((uint32_t) grid, 128, 0, stream); + ggml_cuda_kernel_launch(k_rwkv_rk_fused_f32<128>, launch_params, + (const float *) cur->data, + (const float *) k->data, + (const float *) r->data, + (const float *) v->data, + (const float *) r_k->data, + (float *) dst->data, + C, H); + } +} diff --git a/ggml/src/ggml-cuda/fused-ops.cuh b/ggml/src/ggml-cuda/fused-ops.cuh new file mode 100644 index 000000000000..fd75f30de9d4 --- /dev/null +++ b/ggml/src/ggml-cuda/fused-ops.cuh @@ -0,0 +1,31 @@ +#include "common.cuh" + +void ggml_cuda_op_lerp_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * x_prev, + const ggml_tensor * cur, + const ggml_tensor * weight, + ggml_tensor * dst); + +void ggml_cuda_op_mul_sub_add_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * base, + const ggml_tensor * scale, + const ggml_tensor * value, + ggml_tensor * dst); + +void ggml_cuda_op_add_mul_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * scale, + ggml_tensor * dst); + +void ggml_cuda_op_rwkv_rk_fused( + ggml_backend_cuda_context & ctx, + const ggml_tensor * cur, + const ggml_tensor * k, + const ggml_tensor * r, + const ggml_tensor * v, + const ggml_tensor * r_k, + ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e8aa0673f695..76efb94f1f4b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -25,6 +25,7 @@ #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/fused-ops.cuh" #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" @@ -3705,6 +3706,100 @@ static bool ggml_cuda_should_fuse_add_mul( return true; } +static bool ggml_cuda_should_fuse_rwkv_rk( + const ggml_tensor * mul_kr, + const ggml_tensor * mul_kr_rk, + const ggml_tensor * sum_rows, + const ggml_tensor * mul_v_rk, + const ggml_tensor * reshape, + const ggml_tensor * add, + const ggml_tensor ** cur, + const ggml_tensor ** k, + const ggml_tensor ** r, + const ggml_tensor ** v, + const ggml_tensor ** r_k) { + if (mul_kr->op != GGML_OP_MUL || mul_kr_rk->op != GGML_OP_MUL || + sum_rows->op != GGML_OP_SUM_ROWS || mul_v_rk->op != GGML_OP_MUL || + reshape->op != GGML_OP_RESHAPE || add->op != GGML_OP_ADD) { + return false; + } + + if (sum_rows->src[0] != mul_kr_rk || reshape->src[0] != mul_v_rk) { + return false; + } + + const ggml_tensor * rk_weight = nullptr; + if (mul_kr_rk->src[0] == mul_kr) { + rk_weight = mul_kr_rk->src[1]; + } else if (mul_kr_rk->src[1] == mul_kr) { + rk_weight = mul_kr_rk->src[0]; + } else { + return false; + } + + const ggml_tensor * value = nullptr; + if (mul_v_rk->src[0] == sum_rows) { + value = mul_v_rk->src[1]; + } else if (mul_v_rk->src[1] == sum_rows) { + value = mul_v_rk->src[0]; + } else { + return false; + } + + const ggml_tensor * add_cur = nullptr; + if (add->src[0] == reshape) { + add_cur = add->src[1]; + } else if (add->src[1] == reshape) { + add_cur = add->src[0]; + } else { + return false; + } + + const ggml_tensor * key = mul_kr->src[0]; + const ggml_tensor * rec = mul_kr->src[1]; + if (key->type != GGML_TYPE_F32 || rec->type != GGML_TYPE_F32 || value->type != GGML_TYPE_F32 || + rk_weight->type != GGML_TYPE_F32 || add_cur->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_is_contiguous(key) || !ggml_is_contiguous(rec) || !ggml_is_contiguous(value) || + !ggml_is_contiguous(rk_weight) || !ggml_is_contiguous(add_cur) || !ggml_is_contiguous(add)) { + return false; + } + + if (!ggml_are_same_shape(key, rec) || !ggml_are_same_shape(key, value)) { + return false; + } + + const int64_t head_size = key->ne[0]; + const int64_t heads = key->ne[1]; + const int64_t n_tokens = key->ne[2]; + const int64_t n_embd = head_size * heads; + + if (head_size != 64 && head_size != 128) { + return false; + } + if (rk_weight->ne[0] != head_size || rk_weight->ne[1] != heads || ggml_nelements(rk_weight) != head_size * heads) { + return false; + } + if (sum_rows->ne[0] != 1 || sum_rows->ne[1] != heads || sum_rows->ne[2] != n_tokens) { + return false; + } + if (add->ne[0] != n_embd || add->ne[1] != n_tokens || ggml_nelements(add) != n_embd * n_tokens) { + return false; + } + if (!ggml_are_same_shape(add_cur, add) || ggml_nelements(reshape) != ggml_nelements(add)) { + return false; + } + + *cur = add_cur; + *k = key; + *r = rec; + *v = value; + *r_k = rk_weight; + return true; +} + static int ggml_cuda_find_single_consumer(const ggml_cgraph * cgraph, int node_idx) { if (ggml_node_get_use_count(cgraph, node_idx) != 1) { return -1; @@ -3750,6 +3845,23 @@ static bool ggml_cuda_add_mul_memory_ok( can_read_write_inplace(scale); } +static bool ggml_cuda_rwkv_rk_memory_ok( + const ggml_tensor * cur, + const ggml_tensor * k, + const ggml_tensor * r, + const ggml_tensor * v, + const ggml_tensor * r_k, + const ggml_tensor * dst) { + if (ggml_cuda_tensors_overlap(dst, cur) && !ggml_are_same_layout(dst, cur)) { + return false; + } + + return !ggml_cuda_tensors_overlap(dst, k) && + !ggml_cuda_tensors_overlap(dst, r) && + !ggml_cuda_tensors_overlap(dst, v) && + !ggml_cuda_tensors_overlap(dst, r_k); +} + // returns whether the write (out) nodes overwrite the read nodes in operation static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, @@ -4202,6 +4314,49 @@ static int ggml_cuda_try_fuse( } } + // RWKV r_k correction: + // cur + reshape(v * sum_rows((k * r) * r_k)) + if (node->op == GGML_OP_MUL) { + const int mul_kr_rk_idx = ggml_cuda_find_single_consumer(cgraph, i); + const int sum_rows_idx = mul_kr_rk_idx > i ? ggml_cuda_find_single_consumer(cgraph, mul_kr_rk_idx) : -1; + const int mul_v_rk_idx = sum_rows_idx > i ? ggml_cuda_find_single_consumer(cgraph, sum_rows_idx) : -1; + const int reshape_idx = mul_v_rk_idx > i ? ggml_cuda_find_single_consumer(cgraph, mul_v_rk_idx) : -1; + const int add_idx = reshape_idx > i ? ggml_cuda_find_single_consumer(cgraph, reshape_idx) : -1; + const ggml_tensor * cur = nullptr; + const ggml_tensor * k = nullptr; + const ggml_tensor * r = nullptr; + const ggml_tensor * v = nullptr; + const ggml_tensor * r_k = nullptr; + + const bool chain_ok = add_idx > i && + cgraph->nodes[mul_kr_rk_idx]->op == GGML_OP_MUL && + cgraph->nodes[sum_rows_idx]->op == GGML_OP_SUM_ROWS && + cgraph->nodes[mul_v_rk_idx]->op == GGML_OP_MUL && + cgraph->nodes[reshape_idx]->op == GGML_OP_RESHAPE && + cgraph->nodes[add_idx]->op == GGML_OP_ADD; + + bool range_ok = chain_ok; + for (int j = i + 1; range_ok && j < add_idx; ++j) { + if (j == mul_kr_rk_idx || j == sum_rows_idx || j == mul_v_rk_idx || j == reshape_idx) { + continue; + } + const ggml_tensor * n = cgraph->nodes[j]; + range_ok = ggml_is_empty(n) || n->op == GGML_OP_RESHAPE || n->op == GGML_OP_TRANSPOSE || + n->op == GGML_OP_VIEW || n->op == GGML_OP_PERMUTE || n->op == GGML_OP_NONE || + (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0; + } + + if (range_ok && + ggml_cuda_should_fuse_rwkv_rk(cgraph->nodes[i], cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], + cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], cgraph->nodes[add_idx], + &cur, &k, &r, &v, &r_k)) { + if (ggml_cuda_rwkv_rk_memory_ok(cur, k, r, v, r_k, cgraph->nodes[add_idx])) { + ggml_cuda_op_rwkv_rk_fused(*cuda_ctx, cur, k, r, v, r_k, cgraph->nodes[add_idx]); + return add_idx - i; + } + } + } + // Snake activation: y = x + sin(a*x)^2 * inv_b // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add if (ggml_can_fuse_subgraph(cgraph, i, diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index c5b17d7827f1..510449ac2993 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -66,7 +66,7 @@ static __global__ void rwkv_wkv_f32(const int B, const int T, const int C, const } template -static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * r_k, const float * s, float * dst) { +static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * s, float * dst) { const int tid = threadIdx.x; const int bid = blockIdx.x; @@ -77,7 +77,7 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons const int n_seq_tokens = T / B; float state[head_size]; - __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size], _r_k[head_size]; + __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size]; #ifndef GGML_USE_MUSA #pragma unroll @@ -85,7 +85,6 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons for (int i = 0; i < head_size; i++) { state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i]; } - _r_k[tid] = r_k[head_i * head_size + tid]; for (int t = batch_i * n_seq_tokens * C + head_i * head_size + tid; t < (batch_i + 1) * n_seq_tokens * C + head_i * head_size + tid; t += C) { __syncthreads(); @@ -98,23 +97,15 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons __syncthreads(); float sa = 0; - float rk = 0; #pragma unroll for (int j = 0; j < head_size; j += 4) { - const float4& r = (float4&)(_r[j]); - const float4& k = (float4&)(_k[j]); const float4& kk = (float4&)(_kk[j]); - const float4& r_k = (float4&)(_r_k[j]); const float4& s = (float4&)(state[j]); sa += kk.x * s.x; sa += kk.y * s.y; sa += kk.z * s.z; sa += kk.w * s.w; - rk += r.x * k.x * r_k.x; - rk += r.y * k.y * r_k.y; - rk += r.z * k.z * r_k.z; - rk += r.w * k.w * r_k.w; } sa = -sa; @@ -145,18 +136,17 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons y += s.w * r.w; } dst[t] = y; - dst[T * C + t] = _v * rk; } #pragma unroll for (int i = 0; i < head_size; i++) { - dst[2 * T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; + dst[T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; } } template static __global__ void __launch_bounds__(WARP_SIZE * rows_per_block, 2) -rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * r_k, const float * s, float * dst) { +rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * kk, const float * a, const float * s, float * dst) { constexpr int head_size = CUDA_WKV_BLOCK_SIZE; constexpr int half_head = head_size / 2; @@ -170,7 +160,7 @@ rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r const int head_off = head_i * head_size; const int t = batch_i * C + head_off + row; - __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size], _r_k[head_size]; + __shared__ float _r[head_size], _w[head_size], _k[head_size], _kk[head_size], _a[head_size]; if (threadIdx.y == 0) { constexpr float w_scale = -0.6065306597126334f; @@ -180,14 +170,12 @@ rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r _k[lane] = k[batch_i * C + head_off + lane]; _kk[lane] = kk[batch_i * C + head_off + lane]; _a[lane] = a[batch_i * C + head_off + lane]; - _r_k[lane] = r_k[head_off + lane]; _r[lane + half_head] = r[batch_i * C + head_off + lane + half_head]; _w[lane + half_head] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane + half_head]))); _k[lane + half_head] = k[batch_i * C + head_off + lane + half_head]; _kk[lane + half_head] = kk[batch_i * C + head_off + lane + half_head]; _a[lane + half_head] = a[batch_i * C + head_off + lane + half_head]; - _r_k[lane + half_head] = r_k[head_off + lane + half_head]; } __syncthreads(); @@ -196,20 +184,17 @@ rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r const float s1 = s[state_base + lane + half_head]; const float sa = -warp_reduce_sum(_kk[lane] * s0 + _kk[lane + half_head] * s1); - const float rk = warp_reduce_sum(_r[lane] * _k[lane] * _r_k[lane] + - _r[lane + half_head] * _k[lane + half_head] * _r_k[lane + half_head]); const float vt = v[t]; const float st0 = s0 * _w[lane] + _k[lane] * vt + sa * _kk[lane] * _a[lane]; const float st1 = s1 * _w[lane + half_head] + _k[lane + half_head] * vt + sa * _kk[lane + half_head] * _a[lane + half_head]; const float y = warp_reduce_sum(st0 * _r[lane] + st1 * _r[lane + half_head]); - dst[2 * T * C + state_base + lane] = st0; - dst[2 * T * C + state_base + lane + half_head] = st1; + dst[T * C + state_base + lane] = st0; + dst[T * C + state_base + lane + half_head] = st1; if (lane == 0) { dst[t] = y; - dst[T * C + t] = vt * rk; } } @@ -248,10 +233,9 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const float * v_d = (const float *)dst->src[3]->data; const float * kk_d = (const float *)dst->src[4]->data; const float * a_d = (const float *)dst->src[5]->data; - const float * r_k_d = (const float *)dst->src[6]->data; - const float * s_d = (const float *)dst->src[7]->data; + const float * s_d = (const float *)dst->src[6]->data; - const int64_t B = dst->src[7]->ne[1]; + const int64_t B = dst->src[6]->ne[1]; const int64_t T = dst->src[0]->ne[2]; const int64_t C = dst->ne[0]; const int64_t H = dst->src[0]->ne[1]; @@ -260,16 +244,16 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) cudaStream_t stream = ctx.stream(); - GGML_ASSERT(dst->src[7]->type == GGML_TYPE_F32); + GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2); if (T / B == 1 && C / H == CUDA_WKV_BLOCK_SIZE) { constexpr int rows_per_block = 4; - rwkv_wkv7_f32_t1_warp_row<<>>(T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); + rwkv_wkv7_f32_t1_warp_row<<>>(T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, s_d, dst_d); } else if (C / H == CUDA_WKV_BLOCK_SIZE) { - rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); + rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, s_d, dst_d); } else { - rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d); + rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, s_d, dst_d); } } diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 95b717ac4848..0e1f1de4577d 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -576,7 +576,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_ } break; case GGML_OP_RWKV_WKV7: { - GGML_ASSERT(op->src[7]->type == GGML_TYPE_F32); + GGML_ASSERT(op->src[6]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == 64); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index a5d9d746c0aa..18656b346f21 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1561,7 +1561,7 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[7]->ne[1]; + const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[6]->ne[1]; const int64_t T = op->src[0]->ne[2]; const int64_t C = op->ne[0]; const int64_t H = op->src[0]->ne[1]; @@ -1579,7 +1579,6 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); if (op->op == GGML_OP_RWKV_WKV7) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), ida++); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[7]), ida++); } ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); ggml_metal_encoder_set_bytes (enc, (void *) &B, sizeof(B), ida++); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 96521351e85c..7ea707b615e9 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2482,7 +2482,6 @@ kernel void kernel_rwkv_wkv7_f32( device const float * v, device const float * kk, device const float * a, - device const float * r_k, device const float * state_in, device float * dst, constant uint & B, @@ -2510,7 +2509,6 @@ kernel void kernel_rwkv_wkv7_f32( threadgroup float _k[head_size]; threadgroup float _kk[head_size]; threadgroup float _a[head_size]; - threadgroup float _r_k[head_size]; float state[head_size]; @@ -2518,7 +2516,6 @@ kernel void kernel_rwkv_wkv7_f32( state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + tid * head_size + i]; } - _r_k[tid] = r_k[head_id * head_size + tid]; const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; @@ -2534,22 +2531,16 @@ kernel void kernel_rwkv_wkv7_f32( threadgroup_barrier(mem_flags::mem_threadgroup); const float v_val = v[t]; - float y = 0.0, sa = 0.0, rk = 0.0; + float y = 0.0, sa = 0.0; float4 sa_vec(0.0); - float4 rk_vec(0.0); for (uint j = 0; j < head_size; j += 4) { - float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); - float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); float4 kk_vec = float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); - float4 r_k_vec = float4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]); sa_vec += kk_vec * s_vec; - rk_vec += r_vec * k_vec * r_k_vec; } sa = -(sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3]); - rk = rk_vec[0] + rk_vec[1] + rk_vec[2] + rk_vec[3]; for (uint j = 0; j < head_size; j += 4) { float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); @@ -2571,11 +2562,10 @@ kernel void kernel_rwkv_wkv7_f32( } dst[t] = y; - dst[T * C + t] = v_val * rk; } for (uint i = 0; i < head_size; i++) { - dst[2 * T * C + batch_id * state_size + head_id * head_size * head_size + dst[T * C + batch_id * state_size + head_id * head_size * head_size + tid * head_size + i] = state[i]; } } diff --git a/ggml/src/ggml-sycl/wkv.cpp b/ggml/src/ggml-sycl/wkv.cpp index c8a4832e1059..021c25b6e7c3 100644 --- a/ggml/src/ggml-sycl/wkv.cpp +++ b/ggml/src/ggml-sycl/wkv.cpp @@ -100,7 +100,7 @@ template static void rwkv_wkv7_f32_kernel( const int B, const int T, const int C, const int H, const float* r, const float* w, const float* k, const float* v, - const float* kk, const float* a, const float* r_k, const float* s, + const float* kk, const float* a, const float* s, float* dst, const sycl::nd_item<3>& item_ct1, float* shared_mem) { const int tid = item_ct1.get_local_id(2); @@ -117,7 +117,6 @@ static void rwkv_wkv7_f32_kernel( float* _k = _w + head_size; float* _kk = _k + head_size; float* _a = _kk + head_size; - float* _r_k = _a + head_size; float state[block_size]; @@ -125,7 +124,6 @@ static void rwkv_wkv7_f32_kernel( for (int i = 0; i < head_size; i++) { state[i] = s[batch_i * state_size + head_i * head_size * head_size + tid * head_size + i]; } - _r_k[tid] = r_k[head_i * head_size + tid]; for (int t = batch_i * n_seq_tokens * C + head_i * head_size + tid; t < (batch_i + 1) * n_seq_tokens * C + head_i * head_size + tid; @@ -143,18 +141,15 @@ static void rwkv_wkv7_f32_kernel( item_ct1.barrier(sycl::access::fence_space::local_space); const float _v = v[t]; - float y = 0, sa = 0, rk = 0; - sycl::float4 r4, k4, kk4, r_k4, s4; + float y = 0, sa = 0; + sycl::float4 r4, k4, kk4, s4; #pragma unroll for (int j = 0; j < head_size; j += 4) { r4 = sycl::float4(_r[j], _r[j+1], _r[j+2], _r[j+3]); - k4 = sycl::float4(_k[j], _k[j+1], _k[j+2], _k[j+3]); kk4 = sycl::float4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); - r_k4 = sycl::float4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); s4 = sycl::float4(state[j], state[j+1], state[j+2], state[j+3]); sa += sycl::dot(kk4, s4); - rk += sycl::dot(r4 * k4, r_k4); } sa = -sa; @@ -180,12 +175,11 @@ static void rwkv_wkv7_f32_kernel( } dst[t] = y; - dst[T * C + t] = _v * rk; } #pragma unroll for (int i = 0; i < head_size; i++) { - dst[2 * T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; + dst[T * C + batch_i * state_size + head_i * head_size * head_size + tid * head_size + i] = state[i]; } } @@ -253,23 +247,22 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { const float* v_d = (const float*)dst->src[3]->data; const float* kk_d = (const float*)dst->src[4]->data; const float* a_d = (const float*)dst->src[5]->data; - const float* r_k_d = (const float*)dst->src[6]->data; - const float* s_d = (const float*)dst->src[7]->data; + const float* s_d = (const float*)dst->src[6]->data; float* dst_d = (float*)dst->data; - const int64_t B = dst->src[7]->ne[1]; + const int64_t B = dst->src[6]->ne[1]; const int64_t T = dst->src[0]->ne[2]; const int64_t C = dst->ne[0]; const int64_t H = dst->src[0]->ne[1]; - GGML_ASSERT(dst->src[7]->type == GGML_TYPE_F32); + GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32); GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == WKV_BLOCK_SIZE || C / H == WKV_BLOCK_SIZE * 2); dpct::queue_ptr stream = ctx.stream(); // Calculate execution configuration - const size_t shared_mem_size = C / H * 6 * sizeof(float); // For r, w, k, kk, a, r_k + const size_t shared_mem_size = C / H * 5 * sizeof(float); // For r, w, k, kk, a sycl::range<3> block_dims(1, 1, C / H); sycl::range<3> grid_dims(1, 1, B * H); @@ -282,7 +275,7 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) { rwkv_wkv7_f32_kernel( - B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d, + B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, s_d, dst_d, item_ct1, (float*)shared_mem_acc.get_multi_ptr().get() ); }); @@ -295,7 +288,7 @@ void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) { rwkv_wkv7_f32_kernel( - B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, r_k_d, s_d, dst_d, + B, T, C, H, r_d, w_d, k_d, v_d, kk_d, a_d, s_d, dst_d, item_ct1, (float*)shared_mem_acc.get_multi_ptr().get() ); }); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a2b1ebb93c81..8dcea42d6e5d 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -815,6 +815,7 @@ struct vk_device_struct { vk_pipeline pipeline_mul_norepeat[2][2][2]; vk_pipeline pipeline_add_mul_f32; vk_pipeline pipeline_lerp_f32; + vk_pipeline pipeline_rwkv_rk_f32; vk_pipeline pipeline_div[2][2][2]; vk_pipeline pipeline_div_norepeat[2][2][2]; vk_pipeline pipeline_add_rms[2][2][2]; @@ -1651,6 +1652,18 @@ struct vk_op_lerp_push_constants { uint32_t d_offset; }; +struct vk_op_rwkv_rk_push_constants { + uint32_t C; + uint32_t H; + uint32_t head_size; + uint32_t cur_offset; + uint32_t k_offset; + uint32_t r_offset; + uint32_t v_offset; + uint32_t rk_offset; + uint32_t dst_offset; +}; + struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -2166,6 +2179,8 @@ struct ggml_backend_vk_context { bool fused_add_mul {}; int fused_add_mul_dst_idx {-1}; bool fused_lerp {}; + bool fused_rwkv_rk {}; + int fused_rwkv_rk_dst_idx {-1}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; @@ -5093,6 +5108,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_add_mul_f32, "add_mul_f32", add_mul_f32_len, add_mul_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_lerp_f32, "lerp_f32", lerp_f32_len, lerp_f32_data, "main", 4, sizeof(vk_op_lerp_push_constants), {256, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rwkv_rk_f32, "rwkv_rk_f32", rwkv_rk_f32_len, rwkv_rk_f32_data, "main", 6, sizeof(vk_op_rwkv_rk_push_constants), {128, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true); for (auto &it : device->pipeline_fa_mask_opt) { @@ -5441,10 +5457,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); - ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 9, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); if (device->subgroup_arithmetic && device->subgroup_size <= 64 && 64 % device->subgroup_size == 0) { const uint32_t rows_per_wg = 4; - ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_t1_f32, "rwkv_wkv7_t1_f32", rwkv_wkv7_t1_f32_len, rwkv_wkv7_t1_f32_data, "main", 9, sizeof(vk_op_rwkv_wkv7_push_constants), {1, rows_per_wg, 1}, {device->subgroup_size, rows_per_wg}, 1, true, true, device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_t1_f32, "rwkv_wkv7_t1_f32", rwkv_wkv7_t1_f32_len, rwkv_wkv7_t1_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, rows_per_wg, 1}, {device->subgroup_size, rows_per_wg}, 1, true, true, device->subgroup_size); } { @@ -11976,6 +11992,22 @@ static bool ggml_vk_should_fuse_lerp( const ggml_tensor ** cur, const ggml_tensor ** weight); +static bool ggml_vk_should_fuse_rwkv_rk( + const ggml_backend_vk_context * ctx, + const ggml_tensor * mul_kr, + const ggml_tensor * mul_kr_rk, + const ggml_tensor * sum_rows, + const ggml_tensor * mul_v_rk, + const ggml_tensor * reshape, + const ggml_tensor * add, + const ggml_tensor ** cur, + const ggml_tensor ** k, + const ggml_tensor ** r, + const ggml_tensor ** v, + const ggml_tensor ** r_k); + +static int ggml_vk_find_single_consumer(const ggml_cgraph * cgraph, int node_idx); + static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { const ggml_tensor * sub = cgraph->nodes[node_idx]; const ggml_tensor * repeat = ctx->num_additional_fused_ops == 3 ? cgraph->nodes[node_idx + 1] : nullptr; @@ -12014,6 +12046,62 @@ static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_prev_buf, cur_buf, weight_buf, dst_buf }, pc, { base_total, n_mix, 1 }); } +static void ggml_vk_rwkv_rk_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { + const ggml_tensor * mul_kr = cgraph->nodes[node_idx]; + const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); + const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); + const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); + const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + GGML_ASSERT(mul_kr_rk_idx > node_idx && sum_rows_idx > node_idx && mul_v_rk_idx > node_idx && reshape_idx > node_idx); + const ggml_tensor * mul_kr_rk = cgraph->nodes[mul_kr_rk_idx]; + const ggml_tensor * sum_rows = cgraph->nodes[sum_rows_idx]; + const ggml_tensor * mul_v_rk = cgraph->nodes[mul_v_rk_idx]; + const ggml_tensor * reshape = cgraph->nodes[reshape_idx]; + ggml_tensor * dst = cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]; + + const ggml_tensor * cur = nullptr; + const ggml_tensor * k = nullptr; + const ggml_tensor * r = nullptr; + const ggml_tensor * v = nullptr; + const ggml_tensor * r_k = nullptr; + GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, mul_kr, mul_kr_rk, sum_rows, mul_v_rk, reshape, dst, &cur, &k, &r, &v, &r_k)); + + const uint32_t cur_offset = get_misalign_bytes(ctx, cur) / ggml_type_size(cur->type); + const uint32_t k_offset = get_misalign_bytes(ctx, k) / ggml_type_size(k->type); + const uint32_t r_offset = get_misalign_bytes(ctx, r) / ggml_type_size(r->type); + const uint32_t v_offset = get_misalign_bytes(ctx, v) / ggml_type_size(v->type); + const uint32_t rk_offset = get_misalign_bytes(ctx, r_k) / ggml_type_size(r_k->type); + const uint32_t dst_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + const uint32_t head_size = (uint32_t) k->ne[0]; + const uint32_t H = (uint32_t) k->ne[1]; + const uint32_t T = (uint32_t) k->ne[2]; + const uint32_t C = head_size * H; + + vk_op_rwkv_rk_push_constants pc { + C, + H, + head_size, + cur_offset, + k_offset, + r_offset, + v_offset, + rk_offset, + dst_offset, + }; + + vk_subbuffer cur_buf = ggml_vk_tensor_subbuffer(ctx, cur, true); + vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k, true); + vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, r, true); + vk_subbuffer v_buf = ggml_vk_tensor_subbuffer(ctx, v, true); + vk_subbuffer rk_buf = ggml_vk_tensor_subbuffer(ctx, r_k, true); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_pipeline pipeline = ctx->device->pipeline_rwkv_rk_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { cur_buf, k_buf, r_buf, v_buf, rk_buf, dst_buf }, pc, { 128, H * T, 1 }); +} + static void ggml_vk_div(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); @@ -12046,7 +12134,7 @@ static void ggml_vk_add_id(ggml_backend_vk_context * ctx, vk_context& subctx, co static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_rwkv_wkv6_push_constants&& pc, int version) { GGML_ASSERT(version == 6 || version == 7); - int num_srcs = version == 6 ? 6 : 8; + int num_srcs = version == 6 ? 6 : 7; for (int i = 0; i < num_srcs; i++) { GGML_ASSERT(!ggml_is_quantized(dst->src[i]->type)); @@ -12077,7 +12165,7 @@ static void ggml_vk_op_f32_wkv(ggml_backend_vk_context * ctx, vk_context& subctx pc, elements); } else if (version == 7) { ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], src_buf[7], dst_buf}, + {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], dst_buf}, pc, elements); } else { // shouldn't happen @@ -12091,7 +12179,7 @@ static bool ggml_vk_rwkv_wkv7_t1(ggml_backend_vk_context * ctx, vk_context& subc return false; } - for (int i = 0; i < 8; i++) { + for (int i = 0; i < 7; i++) { GGML_ASSERT(!ggml_is_quantized(dst->src[i]->type)); } @@ -12103,13 +12191,13 @@ static bool ggml_vk_rwkv_wkv7_t1(ggml_backend_vk_context * ctx, vk_context& subc ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); - vk_subbuffer src_buf[8] = {}; - for (int i = 0; i < 8; i++) { + vk_subbuffer src_buf[7] = {}; + for (int i = 0; i < 7; i++) { src_buf[i] = ggml_vk_tensor_subbuffer(ctx, dst->src[i]); } ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], src_buf[7], dst_buf}, + {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], src_buf[6], dst_buf}, pc, {pc.B * pc.H, head_size, 1}); return true; @@ -12137,7 +12225,7 @@ static void ggml_vk_rwkv_wkv7(ggml_backend_vk_context * ctx, vk_context& subctx, const size_t seq_length = dst->src[0]->ne[2]; const size_t n_embed = dst->ne[0]; const size_t n_heads = dst->src[0]->ne[1]; - const size_t n_seqs = dst->src[7]->ne[1]; + const size_t n_seqs = dst->src[6]->ne[1]; const vk_op_rwkv_wkv6_push_constants pc = { (uint32_t)n_seqs, (uint32_t)seq_length, @@ -14756,7 +14844,34 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr return false; }; - if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { + if (ctx->fused_rwkv_rk && ctx->fused_rwkv_rk_dst_idx >= 0) { + const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); + const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); + const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); + const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + + const ggml_tensor * cur = nullptr; + const ggml_tensor * k = nullptr; + const ggml_tensor * r = nullptr; + const ggml_tensor * v = nullptr; + const ggml_tensor * r_k = nullptr; + GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], + cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], + cgraph->nodes[ctx->fused_rwkv_rk_dst_idx], &cur, &k, &r, &v, &r_k)); + + const ggml_tensor * dst = cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]; + const ggml_tensor * fused_srcs[] = { cur, k, r, v, r_k }; + + if (overlaps_unsynced(dst, ctx->unsynced_nodes_read) || overlaps_unsynced(dst, ctx->unsynced_nodes_written)) { + need_sync = true; + } + for (const ggml_tensor * src : fused_srcs) { + if (src && overlaps_unsynced(src, ctx->unsynced_nodes_written)) { + need_sync = true; + break; + } + } + } else if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { const ggml_tensor * add = node; const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; @@ -14809,7 +14924,28 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr ggml_vk_sync_buffers(ctx, compute_ctx); } } - if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { + if (ctx->fused_rwkv_rk && ctx->fused_rwkv_rk_dst_idx >= 0) { + const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); + const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); + const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); + const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + + const ggml_tensor * cur = nullptr; + const ggml_tensor * k = nullptr; + const ggml_tensor * r = nullptr; + const ggml_tensor * v = nullptr; + const ggml_tensor * r_k = nullptr; + GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], + cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], + cgraph->nodes[ctx->fused_rwkv_rk_dst_idx], &cur, &k, &r, &v, &r_k)); + + ctx->unsynced_nodes_written.push_back(cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]); + ctx->unsynced_nodes_read.push_back(cur); + ctx->unsynced_nodes_read.push_back(k); + ctx->unsynced_nodes_read.push_back(r); + ctx->unsynced_nodes_read.push_back(v); + ctx->unsynced_nodes_read.push_back(r_k); + } else if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { const ggml_tensor * add = node; const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; @@ -14889,7 +15025,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_MUL: - if (ctx->num_additional_fused_ops) { + if (ctx->fused_rwkv_rk) { + ggml_vk_rwkv_rk_fused(ctx, compute_ctx, cgraph, node_idx); + } else if (ctx->num_additional_fused_ops) { ggml_vk_snake_dispatch_fused(ctx, compute_ctx, cgraph, node_idx); } else { ggml_vk_mul(ctx, compute_ctx, src0, src1, node); @@ -16065,6 +16203,111 @@ static bool ggml_vk_should_fuse_lerp( return true; } +static bool ggml_vk_should_fuse_rwkv_rk( + const ggml_backend_vk_context * ctx, + const ggml_tensor * mul_kr, + const ggml_tensor * mul_kr_rk, + const ggml_tensor * sum_rows, + const ggml_tensor * mul_v_rk, + const ggml_tensor * reshape, + const ggml_tensor * add, + const ggml_tensor ** cur, + const ggml_tensor ** k, + const ggml_tensor ** r, + const ggml_tensor ** v, + const ggml_tensor ** r_k) { + if (mul_kr->op != GGML_OP_MUL || mul_kr_rk->op != GGML_OP_MUL || + sum_rows->op != GGML_OP_SUM_ROWS || mul_v_rk->op != GGML_OP_MUL || + reshape->op != GGML_OP_RESHAPE || add->op != GGML_OP_ADD) { + return false; + } + + if (sum_rows->src[0] != mul_kr_rk || reshape->src[0] != mul_v_rk) { + return false; + } + + const ggml_tensor * rk_weight = nullptr; + if (mul_kr_rk->src[0] == mul_kr) { + rk_weight = mul_kr_rk->src[1]; + } else if (mul_kr_rk->src[1] == mul_kr) { + rk_weight = mul_kr_rk->src[0]; + } else { + return false; + } + + const ggml_tensor * value = nullptr; + if (mul_v_rk->src[0] == sum_rows) { + value = mul_v_rk->src[1]; + } else if (mul_v_rk->src[1] == sum_rows) { + value = mul_v_rk->src[0]; + } else { + return false; + } + + const ggml_tensor * add_cur = nullptr; + if (add->src[0] == reshape) { + add_cur = add->src[1]; + } else if (add->src[1] == reshape) { + add_cur = add->src[0]; + } else { + return false; + } + + const ggml_tensor * key = mul_kr->src[0]; + const ggml_tensor * rec = mul_kr->src[1]; + if (key->type != GGML_TYPE_F32 || rec->type != GGML_TYPE_F32 || value->type != GGML_TYPE_F32 || + rk_weight->type != GGML_TYPE_F32 || add_cur->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { + return false; + } + + if (!ggml_is_contiguous(key) || !ggml_is_contiguous(rec) || !ggml_is_contiguous(value) || + !ggml_is_contiguous(rk_weight) || !ggml_is_contiguous(add_cur) || !ggml_is_contiguous(add)) { + return false; + } + + if (!ggml_are_same_shape(key, rec) || !ggml_are_same_shape(key, value)) { + return false; + } + + const int64_t head_size = key->ne[0]; + const int64_t heads = key->ne[1]; + const int64_t n_tokens = key->ne[2]; + const int64_t n_embd = head_size * heads; + + if (head_size != 64 && head_size != 128) { + return false; + } + if (rk_weight->ne[0] != head_size || rk_weight->ne[1] != heads || ggml_nelements(rk_weight) != head_size * heads) { + return false; + } + if (sum_rows->ne[0] != 1 || sum_rows->ne[1] != heads || sum_rows->ne[2] != n_tokens) { + return false; + } + if (add->ne[0] != n_embd || add->ne[1] != n_tokens || ggml_nelements(add) != n_embd * n_tokens) { + return false; + } + if (!ggml_are_same_shape(add_cur, add) || ggml_nelements(reshape) != ggml_nelements(add)) { + return false; + } + if (n_embd > UINT32_MAX || heads > UINT32_MAX || n_tokens > UINT32_MAX) { + return false; + } + + const ggml_tensor * tensors[] = { add_cur, key, rec, value, rk_weight, add }; + for (const ggml_tensor * tensor : tensors) { + if (get_misalign_bytes(ctx, tensor) / ggml_type_size(tensor->type) > UINT32_MAX) { + return false; + } + } + + *cur = add_cur; + *k = key; + *r = rec; + *v = value; + *r_k = rk_weight; + return true; +} + static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; @@ -16554,6 +16797,52 @@ static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor !ggml_vk_tensors_overlap(scale, mul, true); } +static bool ggml_vk_can_fuse_rwkv_rk( + const ggml_backend_vk_context * ctx, + const ggml_cgraph * cgraph, + int node_idx, + int * add_idx) { + const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); + const int sum_rows_idx = mul_kr_rk_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx) : -1; + const int mul_v_rk_idx = sum_rows_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, sum_rows_idx) : -1; + const int reshape_idx = mul_v_rk_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx) : -1; + const int found_add_idx = reshape_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, reshape_idx) : -1; + + const bool chain_ok = found_add_idx > node_idx && + cgraph->nodes[node_idx]->op == GGML_OP_MUL && + cgraph->nodes[mul_kr_rk_idx]->op == GGML_OP_MUL && + cgraph->nodes[sum_rows_idx]->op == GGML_OP_SUM_ROWS && + cgraph->nodes[mul_v_rk_idx]->op == GGML_OP_MUL && + cgraph->nodes[reshape_idx]->op == GGML_OP_RESHAPE && + cgraph->nodes[found_add_idx]->op == GGML_OP_ADD; + if (!chain_ok) { + return false; + } + + const ggml_tensor * cur = nullptr; + const ggml_tensor * k = nullptr; + const ggml_tensor * r = nullptr; + const ggml_tensor * v = nullptr; + const ggml_tensor * r_k = nullptr; + if (!ggml_vk_should_fuse_rwkv_rk(ctx, cgraph->nodes[node_idx], cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], + cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], cgraph->nodes[found_add_idx], + &cur, &k, &r, &v, &r_k)) { + return false; + } + + const ggml_tensor * dst = cgraph->nodes[found_add_idx]; + if (ggml_vk_tensors_overlap(cur, dst, true) || + ggml_vk_tensors_overlap(k, dst, true) || + ggml_vk_tensors_overlap(r, dst, true) || + ggml_vk_tensors_overlap(v, dst, true) || + ggml_vk_tensors_overlap(r_k, dst, true)) { + return false; + } + + *add_idx = found_add_idx; + return true; +} + static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { GGML_UNUSED(ctx); @@ -16761,7 +17050,10 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; + ctx->fused_rwkv_rk = false; + ctx->fused_rwkv_rk_dst_idx = -1; const char *fusion_string {}; + int rwkv_rk_add_idx = -1; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); if (num_adds) { @@ -16833,6 +17125,10 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_add_mul_dst_idx = mul_idx; fusion_string = "ADD_MUL"; } + } else if (cgraph->nodes[i]->op == GGML_OP_MUL && ggml_vk_can_fuse_rwkv_rk(ctx, cgraph, i, &rwkv_rk_add_idx)) { + ctx->fused_rwkv_rk = true; + ctx->fused_rwkv_rk_dst_idx = rwkv_rk_add_idx; + fusion_string = "RWKV_RK"; } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { ctx->num_additional_fused_ops = 2; fusion_string = "NORM_MUL_ADD"; @@ -17004,6 +17300,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; + ctx->fused_rwkv_rk = false; + ctx->fused_rwkv_rk_dst_idx = -1; } } @@ -17032,6 +17330,17 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } if (enqueued) { + if (ctx->fused_rwkv_rk_dst_idx >= 0) { + const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, i); + const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); + const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); + const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + fused_noncontiguous_nodes.insert(cgraph->nodes[mul_kr_rk_idx]); + fused_noncontiguous_nodes.insert(cgraph->nodes[sum_rows_idx]); + fused_noncontiguous_nodes.insert(cgraph->nodes[mul_v_rk_idx]); + fused_noncontiguous_nodes.insert(cgraph->nodes[reshape_idx]); + fused_noncontiguous_nodes.insert(cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]); + } if (ctx->fused_add_mul_dst_idx >= 0) { fused_noncontiguous_nodes.insert(cgraph->nodes[ctx->fused_add_mul_dst_idx]); } @@ -17059,6 +17368,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_add_mul = false; ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; + ctx->fused_rwkv_rk = false; + ctx->fused_rwkv_rk_dst_idx = -1; } ctx->last_total_flops = total_flops; @@ -19066,7 +19377,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * src_clone[2], src_clone[3], src_clone[4], src_clone[5]); } else if (tensor->op == GGML_OP_RWKV_WKV7) { tensor_clone = ggml_rwkv_wkv7(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], - src_clone[4], src_clone[5], src_clone[6], src_clone[7]); + src_clone[4], src_clone[5], src_clone[6]); } else if (tensor->op == GGML_OP_GATED_DELTA_NET) { tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rwkv_rk.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rwkv_rk.comp new file mode 100644 index 000000000000..c6746f7126b1 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rwkv_rk.comp @@ -0,0 +1,49 @@ +#version 450 + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint C; + uint H; + uint head_size; + uint cur_offset; + uint k_offset; + uint r_offset; + uint v_offset; + uint rk_offset; + uint dst_offset; +}; + +layout(binding = 0) readonly buffer CurBuf { float cur[]; }; +layout(binding = 1) readonly buffer KBuf { float k[]; }; +layout(binding = 2) readonly buffer RBuf { float r[]; }; +layout(binding = 3) readonly buffer VBuf { float v[]; }; +layout(binding = 4) readonly buffer RKBuf { float r_k[]; }; +layout(binding = 5) buffer DstBuf { float dst[]; }; + +shared float rk[128]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint h = gl_WorkGroupID.y % H; + const uint t = gl_WorkGroupID.y / H; + const uint off = t*C + h*head_size; + + if (tid < head_size) { + rk[tid] = k[k_offset + off + tid] * r[r_offset + off + tid] * r_k[rk_offset + h*head_size + tid]; + } else { + rk[tid] = 0.0f; + } + barrier(); + + for (uint stride = 64; stride > 0; stride >>= 1) { + if (tid < stride) { + rk[tid] += rk[tid + stride]; + } + barrier(); + } + + if (tid < head_size) { + dst[dst_offset + off + tid] = cur[cur_offset + off + tid] + v[v_offset + off + tid] * rk[0]; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index e5208e6b8131..dba429713b70 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -839,6 +839,7 @@ void process_shaders() { string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("lerp_f32", "lerp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); string_to_spv("add_mul_f32", "add_mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + string_to_spv("rwkv_rk_f32", "rwkv_rk.comp", {}); string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp index 851831fd8d61..37f58bc8bd1c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp @@ -18,11 +18,10 @@ layout(binding = 2) readonly buffer KBuf { A_TYPE k[]; }; layout(binding = 3) readonly buffer VBuf { A_TYPE v[]; }; layout(binding = 4) readonly buffer KKBuf { A_TYPE kk[]; }; layout(binding = 5) readonly buffer ABuf { A_TYPE a[]; }; -layout(binding = 6) readonly buffer RKBuf { A_TYPE r_k[]; }; -layout(binding = 7) readonly buffer StateBuf { A_TYPE state_in[]; }; -layout(binding = 8) buffer DstBuf { A_TYPE dst[]; }; +layout(binding = 6) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 7) buffer DstBuf { A_TYPE dst[]; }; -shared A_TYPE _r[BLOCK_SIZE], _w[BLOCK_SIZE], _k[BLOCK_SIZE], _kk[BLOCK_SIZE], _a[BLOCK_SIZE], _r_k[BLOCK_SIZE]; +shared A_TYPE _r[BLOCK_SIZE], _w[BLOCK_SIZE], _k[BLOCK_SIZE], _kk[BLOCK_SIZE], _a[BLOCK_SIZE]; void main() { const uint head_size = BLOCK_SIZE; @@ -42,7 +41,6 @@ void main() { state[i] = state_in[batch_id * state_size + head_id * head_size * head_size + tid * head_size + i]; } - _r_k[tid] = r_k[head_id * head_size + tid]; const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid; const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid; @@ -58,15 +56,10 @@ void main() { barrier(); A_TYPE sa = 0.0; - A_TYPE rk = 0.0; [[unroll]] for (uint j = 0; j < head_size; j += 4) { - vec4 r_vec = vec4(_r[j], _r[j+1], _r[j+2], _r[j+3]); - vec4 k_vec = vec4(_k[j], _k[j+1], _k[j+2], _k[j+3]); vec4 s_vec = vec4(state[j], state[j+1], state[j+2], state[j+3]); vec4 kk_vec = vec4(_kk[j], _kk[j+1], _kk[j+2], _kk[j+3]); - vec4 r_k_vec = vec4(_r_k[j], _r_k[j+1], _r_k[j+2], _r_k[j+3]); sa += dot(s_vec, kk_vec); - rk += dot(r_vec * k_vec, r_k_vec); } sa = -sa; @@ -92,11 +85,10 @@ void main() { } dst[t] = y; - dst[T * C + t] = v_val * rk; } [[unroll]] for (uint i = 0; i < head_size; i++) { - dst[2 * T * C + batch_id * state_size + head_id * head_size * head_size + dst[T * C + batch_id * state_size + head_id * head_size * head_size + tid * head_size + i] = state[i]; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp index 993bed63d8cf..82a519f8050c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp @@ -25,11 +25,10 @@ layout(binding = 2) readonly buffer KBuf { A_TYPE k[]; }; layout(binding = 3) readonly buffer VBuf { A_TYPE v[]; }; layout(binding = 4) readonly buffer KKBuf { A_TYPE kk[]; }; layout(binding = 5) readonly buffer ABuf { A_TYPE a[]; }; -layout(binding = 6) readonly buffer RKBuf { A_TYPE r_k[]; }; -layout(binding = 7) readonly buffer StateBuf { A_TYPE state_in[]; }; -layout(binding = 8) buffer DstBuf { A_TYPE dst[]; }; +layout(binding = 6) readonly buffer StateBuf { A_TYPE state_in[]; }; +layout(binding = 7) buffer DstBuf { A_TYPE dst[]; }; -shared A_TYPE _r[HEAD_SIZE], _w[HEAD_SIZE], _k[HEAD_SIZE], _kk[HEAD_SIZE], _a[HEAD_SIZE], _r_k[HEAD_SIZE]; +shared A_TYPE _r[HEAD_SIZE], _w[HEAD_SIZE], _k[HEAD_SIZE], _kk[HEAD_SIZE], _a[HEAD_SIZE]; void main() { const uint batch_id = gl_WorkGroupID.x / H; @@ -55,7 +54,6 @@ void main() { _k[i] = k[t]; _kk[i] = kk[t]; _a[i] = a[t]; - _r_k[i] = r_k[head_off + i]; } } barrier(); @@ -67,15 +65,12 @@ void main() { } A_TYPE sa_partial = 0.0; - A_TYPE rk_partial = 0.0; [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { const uint i = r_i * SUBGROUP_SIZE + lane; sa_partial += s_shard[r_i] * _kk[i]; - rk_partial += _r[i] * _k[i] * _r_k[i]; } const A_TYPE sa = -subgroupAdd(sa_partial); - const A_TYPE rk = subgroupAdd(rk_partial); const uint t_out = token_base + row; const A_TYPE v_val = v[t_out]; @@ -93,11 +88,10 @@ void main() { if (lane == 0) { dst[t_out] = y; - dst[T * C + t_out] = v_val * rk; } [[unroll]] for (uint r_i = 0; r_i < ROWS_PER_LANE; r_i++) { const uint i = r_i * SUBGROUP_SIZE + lane; - dst[2 * T * C + batch_id * state_size + head_id * HEAD_SIZE * HEAD_SIZE + row * HEAD_SIZE + i] = s_shard[r_i]; + dst[T * C + batch_id * state_size + head_id * HEAD_SIZE * HEAD_SIZE + row * HEAD_SIZE + i] = s_shard[r_i]; } } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index ea3f4b5e84ca..5425a7b5878d 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1169,7 +1169,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "add_rel_pos(x)", "rwkv_wkv6(k, v, r, tf, td, s)", "gated_linear_attn(k, v, q, gate, s)", - "rwkv_wkv7(r, w, k, v, a, b, s)", + "rwkv_wkv7(r, w, k, v, kk, a, s)", "A X = B, A triangular, solve X", "gated_delta_net(q, k, v, g, beta, s)", @@ -5814,7 +5814,6 @@ struct ggml_tensor * ggml_rwkv_wkv7( struct ggml_tensor * v, struct ggml_tensor * kk, struct ggml_tensor * a, - struct ggml_tensor * r_k, struct ggml_tensor * state) { GGML_ASSERT(ggml_is_contiguous(r)); GGML_ASSERT(ggml_is_contiguous(w)); @@ -5822,7 +5821,6 @@ struct ggml_tensor * ggml_rwkv_wkv7( GGML_ASSERT(ggml_is_contiguous(v)); GGML_ASSERT(ggml_is_contiguous(kk)); GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(r_k)); GGML_ASSERT(ggml_is_contiguous(state)); const int64_t S = k->ne[0]; @@ -5835,12 +5833,11 @@ struct ggml_tensor * ggml_rwkv_wkv7( GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); GGML_ASSERT(kk->ne[0] == S && kk->ne[1] == H && kk->ne[2] == n_tokens); GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); - GGML_ASSERT(r_k->ne[0] == S && r_k->ne[1] == H); GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); } - // concat output, v*rk, and new_state - const int64_t ne[4] = { S * H, 2*n_tokens + S * n_seqs, 1, 1 }; + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); result->op = GGML_OP_RWKV_WKV7; @@ -5850,8 +5847,7 @@ struct ggml_tensor * ggml_rwkv_wkv7( result->src[3] = v; result->src[4] = kk; result->src[5] = a; - result->src[6] = r_k; - result->src[7] = state; + result->src[6] = state; return result; } diff --git a/src/models/rwkv7-base.cpp b/src/models/rwkv7-base.cpp index 78d38b821452..50b24e5dccfe 100644 --- a/src/models/rwkv7-base.cpp +++ b/src/models/rwkv7-base.cpp @@ -104,10 +104,9 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in ggml_tensor * wkv_state = build_rs(inp, mctx_cur->get_s_l(il), hparams.n_embd_s(), n_seqs); ggml_tensor * r_k = ggml_reshape_2d(ctx0, layer.time_mix_r_k, head_size, head_count); - ggml_tensor * wkv_output = ggml_rwkv_wkv7(ctx0, r, w, k, v, kk, a, r_k, wkv_state); + ggml_tensor * wkv_output = ggml_rwkv_wkv7(ctx0, r, w, k, v, kk, a, wkv_state); cur = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, 0); - ggml_tensor * vk = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, n_embd * n_tokens * sizeof(float)); - wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, 2 * n_embd * n_tokens * sizeof(float)); + wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float)); ggml_build_forward_expand( gf, ggml_cpy(ctx0, wkv_state, @@ -125,7 +124,8 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in } else { cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens); } - cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, vk, n_embd, n_tokens)); + ggml_tensor * rk = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_mul(ctx0, k, r), r_k)); + cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, ggml_mul(ctx0, v, rk), n_embd, n_tokens)); if (has_gating) { cur = ggml_mul(ctx0, cur, g); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 436a0a3df65e..4cb160cb2e7e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4004,12 +4004,11 @@ struct test_rwkv_wkv7 : public test_case { ggml_tensor * v = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); ggml_tensor * kk = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); ggml_tensor * a = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * r_k = ggml_new_tensor(ctx, type, 2, std::vector{ head_size, head_count }.data()); // Outputs may become NaN with long seqlen without these normalization kk = ggml_l2_norm(ctx, kk, 1e-7F); a = ggml_l2_norm(ctx, a, 1e-7F); ggml_tensor * s = ggml_new_tensor(ctx, type, 2, std::vector{ head_size * head_size * head_count, n_seqs }.data()); - ggml_tensor * out = ggml_rwkv_wkv7(ctx, r, w, k, v, kk, a, r_k, s); + ggml_tensor * out = ggml_rwkv_wkv7(ctx, r, w, k, v, kk, a, s); return out; } }; From 37ed79b25f8007ab8a4b7bbcba6e2ed3d3ab0ef2 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 00:59:23 +0800 Subject: [PATCH 09/16] some style changes Signed-off-by: Molly Sophia --- ggml/src/ggml-cuda/binbcast.cu | 1 - ggml/src/ggml-cuda/wkv.cu | 32 ++++++++++++++++---------------- ggml/src/ggml-sycl/wkv.cpp | 10 +++++----- 3 files changed, 21 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 839c32a48aa7..2e38077bf67f 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -542,7 +542,6 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, } } - void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index 510449ac2993..d7b51a927353 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -165,17 +165,17 @@ rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r if (threadIdx.y == 0) { constexpr float w_scale = -0.6065306597126334f; - _r[lane] = r[batch_i * C + head_off + lane]; - _w[lane] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane]))); - _k[lane] = k[batch_i * C + head_off + lane]; - _kk[lane] = kk[batch_i * C + head_off + lane]; - _a[lane] = a[batch_i * C + head_off + lane]; - - _r[lane + half_head] = r[batch_i * C + head_off + lane + half_head]; - _w[lane + half_head] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane + half_head]))); - _k[lane + half_head] = k[batch_i * C + head_off + lane + half_head]; - _kk[lane + half_head] = kk[batch_i * C + head_off + lane + half_head]; - _a[lane + half_head] = a[batch_i * C + head_off + lane + half_head]; + _r[lane] = r[batch_i * C + head_off + lane]; + _w[lane] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane]))); + _k[lane] = k[batch_i * C + head_off + lane]; + _kk[lane] = kk[batch_i * C + head_off + lane]; + _a[lane] = a[batch_i * C + head_off + lane]; + + _r[lane + half_head] = r[batch_i * C + head_off + lane + half_head]; + _w[lane + half_head] = __expf(w_scale / (1.0f + __expf(-w[batch_i * C + head_off + lane + half_head]))); + _k[lane + half_head] = k[batch_i * C + head_off + lane + half_head]; + _kk[lane + half_head] = kk[batch_i * C + head_off + lane + half_head]; + _a[lane + half_head] = a[batch_i * C + head_off + lane + half_head]; } __syncthreads(); @@ -227,13 +227,13 @@ void ggml_cuda_op_rwkv_wkv6(ggml_backend_cuda_context & ctx, ggml_tensor * dst) } void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const float * r_d = (const float *)dst->src[0]->data; - const float * w_d = (const float *)dst->src[1]->data; - const float * k_d = (const float *)dst->src[2]->data; - const float * v_d = (const float *)dst->src[3]->data; + const float * r_d = (const float *)dst->src[0]->data; + const float * w_d = (const float *)dst->src[1]->data; + const float * k_d = (const float *)dst->src[2]->data; + const float * v_d = (const float *)dst->src[3]->data; const float * kk_d = (const float *)dst->src[4]->data; const float * a_d = (const float *)dst->src[5]->data; - const float * s_d = (const float *)dst->src[6]->data; + const float * s_d = (const float *)dst->src[6]->data; const int64_t B = dst->src[6]->ne[1]; const int64_t T = dst->src[0]->ne[2]; diff --git a/ggml/src/ggml-sycl/wkv.cpp b/ggml/src/ggml-sycl/wkv.cpp index 021c25b6e7c3..bccbbc71fdaf 100644 --- a/ggml/src/ggml-sycl/wkv.cpp +++ b/ggml/src/ggml-sycl/wkv.cpp @@ -241,13 +241,13 @@ void ggml_sycl_op_rwkv_wkv6(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/7); - const float* r_d = (const float*)dst->src[0]->data; - const float* w_d = (const float*)dst->src[1]->data; - const float* k_d = (const float*)dst->src[2]->data; - const float* v_d = (const float*)dst->src[3]->data; + const float* r_d = (const float*)dst->src[0]->data; + const float* w_d = (const float*)dst->src[1]->data; + const float* k_d = (const float*)dst->src[2]->data; + const float* v_d = (const float*)dst->src[3]->data; const float* kk_d = (const float*)dst->src[4]->data; const float* a_d = (const float*)dst->src[5]->data; - const float* s_d = (const float*)dst->src[6]->data; + const float* s_d = (const float*)dst->src[6]->data; float* dst_d = (float*)dst->data; const int64_t B = dst->src[6]->ne[1]; From ca6d97250fc0b17d5ca4c7c233a1a588341ebe66 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 13:06:31 +0800 Subject: [PATCH 10/16] Reduce RWKV fusion scheduler intrusiveness --- ggml/src/ggml-cpu/ggml-cpu.c | 52 +++----------- ggml/src/ggml-cuda/ggml-cuda.cu | 37 +++------- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 100 ++++++++------------------- 3 files changed, 48 insertions(+), 141 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 3ec719a09aea..82f8ecb01070 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2981,30 +2981,11 @@ struct ggml_cplan ggml_graph_plan( // Returns the number of nodes skipped by fusion (>=1), or 0 if no fusion was applied. static bool ggml_cpu_disable_fusion = false; // initialized once in ggml_cpu_init(), read-only afterwards -static int ggml_cpu_find_single_consumer(const struct ggml_cgraph * cgraph, const int node_n) { - if (ggml_node_get_use_count(cgraph, node_n) != 1) { - return -1; - } - - const struct ggml_tensor * node = cgraph->nodes[node_n]; - for (int i = node_n + 1; i < cgraph->n_nodes; ++i) { - const struct ggml_tensor * consumer = cgraph->nodes[i]; - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - if (consumer->src[src_idx] == node) { - return i; - } - } - } - - return -1; -} - static int ggml_cpu_try_fuse_ops( const struct ggml_cgraph * cgraph, const int node_n, const struct ggml_compute_params * params, - const struct ggml_cplan * cplan, - int * skip_node_n) { + const struct ggml_cplan * cplan) { if (ggml_cpu_disable_fusion || cplan->use_ref) { return 0; @@ -3014,9 +2995,9 @@ static int ggml_cpu_try_fuse_ops( if (node->op == GGML_OP_ADD) { // ADD + MUL fusion - const int mul_node_n = ggml_cpu_find_single_consumer(cgraph, node_n); - if (mul_node_n > node_n && cgraph->nodes[mul_node_n]->op == GGML_OP_MUL) { - struct ggml_tensor * mul_node = cgraph->nodes[mul_node_n]; + const enum ggml_op fuse_ops[] = { GGML_OP_ADD, GGML_OP_MUL }; + if (ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) { + struct ggml_tensor * mul_node = cgraph->nodes[node_n + 1]; const struct ggml_tensor * scale = (mul_node->src[0] == node) ? mul_node->src[1] : mul_node->src[0]; @@ -3034,11 +3015,7 @@ static int ggml_cpu_try_fuse_ops( ggml_is_contiguous(mul_node)) { ggml_compute_forward_add_mul_fused(params, node, mul_node); - if (mul_node_n == node_n + 1) { - return 1; - } - *skip_node_n = mul_node_n; - return -1; + return 1; } } } @@ -3124,7 +3101,6 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { GGML_PRINT_DEBUG("thread #%d compute-start cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); #endif - int skip_node_n = -1; for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) { struct ggml_tensor * node = cgraph->nodes[node_n]; @@ -3137,19 +3113,13 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { continue; } - if (skip_node_n == node_n) { - skip_node_n = -1; + // TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time + // Try fused ops, fall back to normal compute + const int n_fused = ggml_cpu_try_fuse_ops(cgraph, node_n, ¶ms, cplan); + if (n_fused > 0) { + node_n += n_fused; } else { - // TODO: move fused-op detection into ggml_graph_plan so fusion decisions are made once at planning time - // Try fused ops, fall back to normal compute - const int n_fused = ggml_cpu_try_fuse_ops(cgraph, node_n, ¶ms, cplan, &skip_node_n); - if (n_fused != 0) { - if (n_fused > 0) { - node_n += n_fused; - } - } else { - ggml_compute_forward(¶ms, node); - } + ggml_compute_forward(¶ms, node); } if (state->ith == 0 && cplan->abort_callback && diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 76efb94f1f4b..99323b355282 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -87,7 +87,6 @@ #include #include #include -#include #include static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); @@ -4150,13 +4149,11 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } // Try and fuse nodes. A positive return value is the number of following -// contiguous nodes to skip; -1 means the current node was handled but no -// following nodes can be skipped because the fused consumer is non-contiguous. +// contiguous nodes to skip. static int ggml_cuda_try_fuse( ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, - int i, - std::unordered_set & fused_noncontiguous_nodes) { + int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { @@ -4292,25 +4289,17 @@ static int ggml_cuda_try_fuse( } } - // Elementwise gate: (a + b) * scale. RWKV builds this as ADD consumed by - // a later MUL, with unrelated view nodes often scheduled between them. - if (node->op == GGML_OP_ADD) { - const int mul_idx = ggml_cuda_find_single_consumer(cgraph, i); + // Elementwise gate: (a + b) * scale. + if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_ADD, GGML_OP_MUL }, { i + 1 })) { const ggml_tensor * src0 = nullptr; const ggml_tensor * src1 = nullptr; const ggml_tensor * scale = nullptr; - const bool should_fuse = mul_idx > i && cgraph->nodes[mul_idx]->op == GGML_OP_MUL && - ggml_cuda_should_fuse_add_mul(cgraph->nodes[i], cgraph->nodes[mul_idx], &src0, &src1, &scale); - const bool memory_ok = should_fuse && ggml_cuda_add_mul_memory_ok(src0, src1, scale, cgraph->nodes[mul_idx]); - - if (memory_ok) { - ggml_cuda_op_add_mul_fused(*cuda_ctx, src0, src1, scale, cgraph->nodes[mul_idx]); - if (mul_idx == i + 1) { + if (ggml_cuda_should_fuse_add_mul(cgraph->nodes[i], cgraph->nodes[i + 1], &src0, &src1, &scale)) { + if (ggml_cuda_add_mul_memory_ok(src0, src1, scale, cgraph->nodes[i + 1])) { + ggml_cuda_op_add_mul_fused(*cuda_ctx, src0, src1, scale, cgraph->nodes[i + 1]); return 1; } - fused_noncontiguous_nodes.insert(cgraph->nodes[mul_idx]); - return -1; } } @@ -4776,8 +4765,6 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud stream_ctx.concurrent_events.clear(); } - std::unordered_set fused_noncontiguous_nodes; - for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -4826,16 +4813,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } - if (fused_noncontiguous_nodes.erase(node) > 0) { - continue; - } - - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, fused_noncontiguous_nodes); + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); if (nodes_to_skip != 0) { - if (nodes_to_skip > 0) { - i += nodes_to_skip; - } + i += nodes_to_skip; continue; } #ifndef NDEBUG diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8dcea42d6e5d..70fae6302e85 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -62,7 +62,6 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #include #include #include -#include #include #include #include @@ -2177,10 +2176,8 @@ struct ggml_backend_vk_context { // If there's no fusion, bit 0 is still set. int fused_ops_write_mask {}; bool fused_add_mul {}; - int fused_add_mul_dst_idx {-1}; bool fused_lerp {}; bool fused_rwkv_rk {}; - int fused_rwkv_rk_dst_idx {-1}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; @@ -11940,8 +11937,7 @@ static void ggml_vk_mul(ggml_backend_vk_context * ctx, vk_context& subctx, const static void ggml_vk_add_mul(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { const ggml_tensor * add = cgraph->nodes[node_idx]; - const int dst_idx = ctx->fused_add_mul_dst_idx >= 0 ? ctx->fused_add_mul_dst_idx : node_idx + 1; - ggml_tensor * dst = cgraph->nodes[dst_idx]; + ggml_tensor * dst = cgraph->nodes[node_idx + 1]; const ggml_tensor * scale = dst->src[0] == add ? dst->src[1] : dst->src[0]; const ggml_tensor * src0 = add->src[0]; const ggml_tensor * src1 = add->src[1]; @@ -12006,6 +12002,8 @@ static bool ggml_vk_should_fuse_rwkv_rk( const ggml_tensor ** v, const ggml_tensor ** r_k); +static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor * mul); + static int ggml_vk_find_single_consumer(const ggml_cgraph * cgraph, int node_idx); static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { @@ -12057,7 +12055,7 @@ static void ggml_vk_rwkv_rk_fused(ggml_backend_vk_context * ctx, vk_context& sub const ggml_tensor * sum_rows = cgraph->nodes[sum_rows_idx]; const ggml_tensor * mul_v_rk = cgraph->nodes[mul_v_rk_idx]; const ggml_tensor * reshape = cgraph->nodes[reshape_idx]; - ggml_tensor * dst = cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]; + ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; const ggml_tensor * cur = nullptr; const ggml_tensor * k = nullptr; @@ -14844,11 +14842,12 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr return false; }; - if (ctx->fused_rwkv_rk && ctx->fused_rwkv_rk_dst_idx >= 0) { + if (ctx->fused_rwkv_rk) { const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + const int add_idx = node_idx + ctx->num_additional_fused_ops; const ggml_tensor * cur = nullptr; const ggml_tensor * k = nullptr; @@ -14857,9 +14856,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr const ggml_tensor * r_k = nullptr; GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], - cgraph->nodes[ctx->fused_rwkv_rk_dst_idx], &cur, &k, &r, &v, &r_k)); + cgraph->nodes[add_idx], &cur, &k, &r, &v, &r_k)); - const ggml_tensor * dst = cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]; + const ggml_tensor * dst = cgraph->nodes[add_idx]; const ggml_tensor * fused_srcs[] = { cur, k, r, v, r_k }; if (overlaps_unsynced(dst, ctx->unsynced_nodes_read) || overlaps_unsynced(dst, ctx->unsynced_nodes_written)) { @@ -14871,21 +14870,6 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; } } - } else if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { - const ggml_tensor * add = node; - const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; - const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; - const ggml_tensor * fused_srcs[] = { add->src[0], add->src[1], scale }; - - if (overlaps_unsynced(mul, ctx->unsynced_nodes_read) || overlaps_unsynced(mul, ctx->unsynced_nodes_written)) { - need_sync = true; - } - for (const ggml_tensor * src : fused_srcs) { - if (src && overlaps_unsynced(src, ctx->unsynced_nodes_written)) { - need_sync = true; - break; - } - } } else { // For all fused ops, check if the destination node or any of the source // nodes require synchronization. @@ -14924,11 +14908,12 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr ggml_vk_sync_buffers(ctx, compute_ctx); } } - if (ctx->fused_rwkv_rk && ctx->fused_rwkv_rk_dst_idx >= 0) { + if (ctx->fused_rwkv_rk) { const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); + const int add_idx = node_idx + ctx->num_additional_fused_ops; const ggml_tensor * cur = nullptr; const ggml_tensor * k = nullptr; @@ -14937,22 +14922,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr const ggml_tensor * r_k = nullptr; GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], - cgraph->nodes[ctx->fused_rwkv_rk_dst_idx], &cur, &k, &r, &v, &r_k)); + cgraph->nodes[add_idx], &cur, &k, &r, &v, &r_k)); - ctx->unsynced_nodes_written.push_back(cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]); + ctx->unsynced_nodes_written.push_back(cgraph->nodes[add_idx]); ctx->unsynced_nodes_read.push_back(cur); ctx->unsynced_nodes_read.push_back(k); ctx->unsynced_nodes_read.push_back(r); ctx->unsynced_nodes_read.push_back(v); ctx->unsynced_nodes_read.push_back(r_k); - } else if (ctx->fused_add_mul && ctx->fused_add_mul_dst_idx >= 0) { - const ggml_tensor * add = node; - const ggml_tensor * mul = cgraph->nodes[ctx->fused_add_mul_dst_idx]; - const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; - ctx->unsynced_nodes_written.push_back(mul); - ctx->unsynced_nodes_read.push_back(add->src[0]); - ctx->unsynced_nodes_read.push_back(add->src[1]); - ctx->unsynced_nodes_read.push_back(scale); } else { // Add all fused nodes to the unsynchronized lists. for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { @@ -16348,7 +16325,8 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g return false; } } else if (ops.size() == 2 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_MUL) { - if (!ggml_vk_should_fuse_add_mul(ctx, cgraph->nodes[node_idx], cgraph->nodes[node_idx + 1])) { + if (!ggml_vk_should_fuse_add_mul(ctx, cgraph->nodes[node_idx], cgraph->nodes[node_idx + 1]) || + !ggml_vk_add_mul_memory_ok(cgraph->nodes[node_idx], cgraph->nodes[node_idx + 1])) { return false; } } else if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { @@ -16819,6 +16797,18 @@ static bool ggml_vk_can_fuse_rwkv_rk( return false; } + for (int i = node_idx + 1; i < found_add_idx; ++i) { + if (i == mul_kr_rk_idx || i == sum_rows_idx || i == mul_v_rk_idx || i == reshape_idx) { + continue; + } + const ggml_tensor * node = cgraph->nodes[i]; + if (!ggml_is_empty(node) && node->op != GGML_OP_RESHAPE && node->op != GGML_OP_TRANSPOSE && + node->op != GGML_OP_VIEW && node->op != GGML_OP_PERMUTE && node->op != GGML_OP_NONE && + (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0) { + return false; + } + } + const ggml_tensor * cur = nullptr; const ggml_tensor * k = nullptr; const ggml_tensor * r = nullptr; @@ -17023,16 +17013,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg uint64_t batch_flops = 0; uint64_t total_flops = 0; uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u); - std::unordered_set fused_noncontiguous_nodes; for (int i = 0; i < cgraph->n_nodes; i++) { if (first_node_in_batch) { submit_node_idx = i; } - if (fused_noncontiguous_nodes.erase(cgraph->nodes[i]) > 0) { - continue; - } - { auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]); batch_flops += node_flops; @@ -17048,10 +17033,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; - ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; ctx->fused_rwkv_rk = false; - ctx->fused_rwkv_rk_dst_idx = -1; const char *fusion_string {}; int rwkv_rk_add_idx = -1; if (!ctx->device->disable_fusion) { @@ -17115,20 +17098,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "ADD_MUL"; op_srcs_fused_elementwise[0] = true; op_srcs_fused_elementwise[1] = true; - } else if (cgraph->nodes[i]->op == GGML_OP_ADD) { - const int mul_idx = ggml_vk_find_single_consumer(cgraph, i); - if (mul_idx > i && - cgraph->nodes[mul_idx]->op == GGML_OP_MUL && - ggml_vk_should_fuse_add_mul(ctx, cgraph->nodes[i], cgraph->nodes[mul_idx]) && - ggml_vk_add_mul_memory_ok(cgraph->nodes[i], cgraph->nodes[mul_idx])) { - ctx->fused_add_mul = true; - ctx->fused_add_mul_dst_idx = mul_idx; - fusion_string = "ADD_MUL"; - } } else if (cgraph->nodes[i]->op == GGML_OP_MUL && ggml_vk_can_fuse_rwkv_rk(ctx, cgraph, i, &rwkv_rk_add_idx)) { ctx->fused_rwkv_rk = true; - ctx->fused_rwkv_rk_dst_idx = rwkv_rk_add_idx; + ctx->num_additional_fused_ops = rwkv_rk_add_idx - i; fusion_string = "RWKV_RK"; + std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false); } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { ctx->num_additional_fused_ops = 2; fusion_string = "NORM_MUL_ADD"; @@ -17298,10 +17272,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; - ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; ctx->fused_rwkv_rk = false; - ctx->fused_rwkv_rk_dst_idx = -1; } } @@ -17330,20 +17302,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } if (enqueued) { - if (ctx->fused_rwkv_rk_dst_idx >= 0) { - const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, i); - const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); - const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); - const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); - fused_noncontiguous_nodes.insert(cgraph->nodes[mul_kr_rk_idx]); - fused_noncontiguous_nodes.insert(cgraph->nodes[sum_rows_idx]); - fused_noncontiguous_nodes.insert(cgraph->nodes[mul_v_rk_idx]); - fused_noncontiguous_nodes.insert(cgraph->nodes[reshape_idx]); - fused_noncontiguous_nodes.insert(cgraph->nodes[ctx->fused_rwkv_rk_dst_idx]); - } - if (ctx->fused_add_mul_dst_idx >= 0) { - fused_noncontiguous_nodes.insert(cgraph->nodes[ctx->fused_add_mul_dst_idx]); - } ++submitted_nodes; #ifndef GGML_VULKAN_CHECK_RESULTS @@ -17366,10 +17324,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 0; ctx->fused_add_mul = false; - ctx->fused_add_mul_dst_idx = -1; ctx->fused_lerp = false; ctx->fused_rwkv_rk = false; - ctx->fused_rwkv_rk_dst_idx = -1; } ctx->last_total_flops = total_flops; From 7029af72bfc61894f01f3903779911f1f2c3fbb9 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 21:20:04 +0800 Subject: [PATCH 11/16] rwkv: use explicit ops for lerp and rk correction --- ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 16 + ggml/src/ggml-backend-meta.cpp | 2 + ggml/src/ggml-cpu/ggml-cpu.c | 13 + ggml/src/ggml-cpu/ops.cpp | 154 +++++++ ggml/src/ggml-cpu/ops.h | 2 + ggml/src/ggml-cuda/ggml-cuda.cu | 275 +---------- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 660 ++++++++------------------- ggml/src/ggml.c | 93 +++- src/models/rwkv7-base.cpp | 9 +- 10 files changed, 469 insertions(+), 759 deletions(-) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5ad121ae57f1..34fad5fb7073 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 1 +#define RPC_PROTO_PATCH_VERSION 2 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 4f96588404ae..77e0a2c0eff4 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -563,6 +563,8 @@ extern "C" { GGML_OP_WIN_UNPART, GGML_OP_GET_REL_POS, GGML_OP_ADD_REL_POS, + GGML_OP_RWKV_LERP, + GGML_OP_RWKV_RK, GGML_OP_RWKV_WKV6, GGML_OP_GATED_LINEAR_ATTN, GGML_OP_RWKV_WKV7, @@ -2501,6 +2503,20 @@ extern "C" { struct ggml_tensor * pw, struct ggml_tensor * ph); + GGML_API struct ggml_tensor * ggml_rwkv_lerp( + struct ggml_context * ctx, + struct ggml_tensor * x_prev, + struct ggml_tensor * cur, + struct ggml_tensor * weight); + + GGML_API struct ggml_tensor * ggml_rwkv_rk( + struct ggml_context * ctx, + struct ggml_tensor * cur, + struct ggml_tensor * k, + struct ggml_tensor * r, + struct ggml_tensor * v, + struct ggml_tensor * r_k); + GGML_API struct ggml_tensor * ggml_rwkv_wkv6( struct ggml_context * ctx, struct ggml_tensor * k, diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 0a36f099000f..fcb88afd2fda 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -975,6 +975,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_OP_WIN_UNPART: case GGML_OP_GET_REL_POS: case GGML_OP_ADD_REL_POS: + case GGML_OP_RWKV_LERP: + case GGML_OP_RWKV_RK: case GGML_OP_RWKV_WKV6: case GGML_OP_GATED_LINEAR_ATTN: case GGML_OP_RWKV_WKV7: diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 82f8ecb01070..f60d142ef058 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2031,6 +2031,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_add_rel_pos(params, tensor); } break; + case GGML_OP_RWKV_LERP: + { + ggml_compute_forward_rwkv_lerp(params, tensor); + } break; + case GGML_OP_RWKV_RK: + { + ggml_compute_forward_rwkv_rk(params, tensor); + } break; case GGML_OP_RWKV_WKV6: { ggml_compute_forward_rwkv_wkv6(params, tensor); @@ -2374,6 +2382,11 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { { n_tasks = n_threads; } break; + case GGML_OP_RWKV_LERP: + case GGML_OP_RWKV_RK: + { + n_tasks = n_threads; + } break; case GGML_OP_RWKV_WKV6: case GGML_OP_GATED_LINEAR_ATTN: case GGML_OP_RWKV_WKV7: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index bc60564ab281..ac44faf5a8f8 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10263,6 +10263,160 @@ void ggml_compute_forward_add_rel_pos( } } +// ggml_compute_forward_rwkv_lerp + +static void ggml_compute_forward_rwkv_lerp_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * x_prev = dst->src[0]; + const ggml_tensor * cur = dst->src[1]; + const ggml_tensor * weight = dst->src[2]; + + GGML_ASSERT(ggml_is_contiguous(x_prev)); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(weight)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t n = ggml_nelements(dst); + const int64_t dr = (n + nth - 1) / nth; + const int64_t i0 = dr * ith; + const int64_t i1 = MIN(i0 + dr, n); + + const float * x = (const float *) x_prev->data; + const float * c = (const float *) cur->data; + const float * w = (const float *) weight->data; + float * d = (float *) dst->data; + + const int64_t base_total = ggml_nelements(x_prev); + const int64_t n_mix = dst->ne[3]; + if (ggml_are_same_shape(x_prev, cur) && + dst->ne[0] == x_prev->ne[0] && + dst->ne[1] == x_prev->ne[1] && + dst->ne[2] == x_prev->ne[2] && + n == base_total * n_mix && + weight->ne[0] == dst->ne[0] && + weight->ne[1] == 1 && + weight->ne[2] == 1 && + weight->ne[3] == n_mix && + ggml_nelements(weight) == dst->ne[0] * n_mix) { + for (int64_t i = i0; i < i1; ++i) { + const int64_t ibase = i % base_total; + const int64_t imix = i / base_total; + const int64_t iembd = ibase % dst->ne[0]; + const float cv = c[ibase]; + d[i] = cv + (x[ibase] - cv) * w[imix * dst->ne[0] + iembd]; + } + return; + } + + const int64_t ne0 = dst->ne[0]; + const int64_t ne1 = dst->ne[1]; + const int64_t ne2 = dst->ne[2]; + + const auto index = [](const ggml_tensor * t, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + const int64_t j0 = i0 % t->ne[0]; + const int64_t j1 = i1 % t->ne[1]; + const int64_t j2 = i2 % t->ne[2]; + const int64_t j3 = i3 % t->ne[3]; + return (((size_t) j3 * t->ne[2] + j2) * t->ne[1] + j1) * t->ne[0] + j0; + }; + + for (int64_t i = i0; i < i1; ++i) { + const int64_t d0 = i % ne0; + const int64_t q1 = i / ne0; + const int64_t d1 = q1 % ne1; + const int64_t q2 = q1 / ne1; + const int64_t d2 = q2 % ne2; + const int64_t d3 = q2 / ne2; + + const size_t ix = index(x_prev, d0, d1, d2, d3); + const size_t ic = index(cur, d0, d1, d2, d3); + const size_t iw = index(weight, d0, d1, d2, d3); + + const float cv = c[ic]; + d[i] = cv + (x[ix] - cv) * w[iw]; + } +} + +void ggml_compute_forward_rwkv_lerp( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + ggml_compute_forward_rwkv_lerp_f32(params, dst); + break; + default: + GGML_ABORT("fatal error"); + } +} + +// ggml_compute_forward_rwkv_rk + +static void ggml_compute_forward_rwkv_rk_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * cur = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * r = dst->src[2]; + const ggml_tensor * v = dst->src[3]; + const ggml_tensor * r_k = dst->src[4]; + + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r_k)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t head_size = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t T = k->ne[2]; + const int64_t C = head_size * H; + + const int ith = params->ith; + const int nth = params->nth; + const int64_t n_rows = H * T; + const int64_t dr = (n_rows + nth - 1) / nth; + const int64_t ir0 = dr * ith; + const int64_t ir1 = MIN(ir0 + dr, n_rows); + + const float * cur_data = (const float *) cur->data; + const float * k_data = (const float *) k->data; + const float * r_data = (const float *) r->data; + const float * v_data = (const float *) v->data; + const float * rk_data = (const float *) r_k->data; + float * dst_data = (float *) dst->data; + + for (int64_t row = ir0; row < ir1; ++row) { + const int64_t h = row % H; + const int64_t t = row / H; + const int64_t off = t * C + h * head_size; + + float rk = 0.0f; + for (int64_t i = 0; i < head_size; ++i) { + rk += k_data[off + i] * r_data[off + i] * rk_data[h * head_size + i]; + } + for (int64_t i = 0; i < head_size; ++i) { + dst_data[off + i] = cur_data[off + i] + v_data[off + i] * rk; + } + } +} + +void ggml_compute_forward_rwkv_rk( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + ggml_compute_forward_rwkv_rk_f32(params, dst); + break; + default: + GGML_ABORT("fatal error"); + } +} + // ggml_compute_forward_rwkv_wkv6 static void ggml_compute_forward_rwkv_wkv6_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index efc8004ce551..21c087e118be 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -102,6 +102,8 @@ void ggml_compute_forward_unary(const struct ggml_compute_params * params, struc void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_lerp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_rk(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 99323b355282..2ad48b05e51a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3093,6 +3093,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_TRI: ggml_cuda_op_tri(ctx, dst); break; + case GGML_OP_RWKV_LERP: + ggml_cuda_op_lerp_fused(ctx, dst->src[0], dst->src[1], dst->src[2], dst); + break; + case GGML_OP_RWKV_RK: + ggml_cuda_op_rwkv_rk_fused(ctx, dst->src[0], dst->src[1], dst->src[2], dst->src[3], dst->src[4], dst); + break; case GGML_OP_RWKV_WKV6: ggml_cuda_op_rwkv_wkv6(ctx, dst); break; @@ -3551,69 +3557,6 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod return true; } -static bool ggml_cuda_should_fuse_lerp( - const ggml_tensor * sub, - const ggml_tensor * repeat, - const ggml_tensor * mul, - const ggml_tensor * add, - const ggml_tensor ** x_prev, - const ggml_tensor ** cur, - const ggml_tensor ** weight) { - if (sub->op != GGML_OP_SUB || mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD) { - return false; - } - - const ggml_tensor * sx = sub; - if (repeat) { - if (repeat->op != GGML_OP_REPEAT || repeat->src[0] != sub) { - return false; - } - sx = repeat; - } - - const ggml_tensor * w = nullptr; - if (mul->src[0] == sx) { - w = mul->src[1]; - } else if (mul->src[1] == sx) { - w = mul->src[0]; - } else { - return false; - } - - const ggml_tensor * add_cur = nullptr; - if (add->src[0] == mul) { - add_cur = add->src[1]; - } else if (add->src[1] == mul) { - add_cur = add->src[0]; - } else { - return false; - } - - if (add_cur != sub->src[1]) { - return false; - } - - const ggml_tensor * xp = sub->src[0]; - const ggml_tensor * c = sub->src[1]; - - if (xp->type != GGML_TYPE_F32 || c->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { - return false; - } - - if (!ggml_can_repeat(xp, add) || !ggml_can_repeat(c, add) || !ggml_can_repeat(w, add)) { - return false; - } - - if (!ggml_is_contiguous(xp) || !ggml_is_contiguous(c) || !ggml_is_contiguous(w) || !ggml_is_contiguous(add)) { - return false; - } - - *x_prev = xp; - *cur = c; - *weight = w; - return true; -} - static bool ggml_cuda_should_fuse_mul_sub_add( const ggml_tensor * mul, const ggml_tensor * sub, @@ -3705,118 +3648,6 @@ static bool ggml_cuda_should_fuse_add_mul( return true; } -static bool ggml_cuda_should_fuse_rwkv_rk( - const ggml_tensor * mul_kr, - const ggml_tensor * mul_kr_rk, - const ggml_tensor * sum_rows, - const ggml_tensor * mul_v_rk, - const ggml_tensor * reshape, - const ggml_tensor * add, - const ggml_tensor ** cur, - const ggml_tensor ** k, - const ggml_tensor ** r, - const ggml_tensor ** v, - const ggml_tensor ** r_k) { - if (mul_kr->op != GGML_OP_MUL || mul_kr_rk->op != GGML_OP_MUL || - sum_rows->op != GGML_OP_SUM_ROWS || mul_v_rk->op != GGML_OP_MUL || - reshape->op != GGML_OP_RESHAPE || add->op != GGML_OP_ADD) { - return false; - } - - if (sum_rows->src[0] != mul_kr_rk || reshape->src[0] != mul_v_rk) { - return false; - } - - const ggml_tensor * rk_weight = nullptr; - if (mul_kr_rk->src[0] == mul_kr) { - rk_weight = mul_kr_rk->src[1]; - } else if (mul_kr_rk->src[1] == mul_kr) { - rk_weight = mul_kr_rk->src[0]; - } else { - return false; - } - - const ggml_tensor * value = nullptr; - if (mul_v_rk->src[0] == sum_rows) { - value = mul_v_rk->src[1]; - } else if (mul_v_rk->src[1] == sum_rows) { - value = mul_v_rk->src[0]; - } else { - return false; - } - - const ggml_tensor * add_cur = nullptr; - if (add->src[0] == reshape) { - add_cur = add->src[1]; - } else if (add->src[1] == reshape) { - add_cur = add->src[0]; - } else { - return false; - } - - const ggml_tensor * key = mul_kr->src[0]; - const ggml_tensor * rec = mul_kr->src[1]; - if (key->type != GGML_TYPE_F32 || rec->type != GGML_TYPE_F32 || value->type != GGML_TYPE_F32 || - rk_weight->type != GGML_TYPE_F32 || add_cur->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { - return false; - } - - if (!ggml_is_contiguous(key) || !ggml_is_contiguous(rec) || !ggml_is_contiguous(value) || - !ggml_is_contiguous(rk_weight) || !ggml_is_contiguous(add_cur) || !ggml_is_contiguous(add)) { - return false; - } - - if (!ggml_are_same_shape(key, rec) || !ggml_are_same_shape(key, value)) { - return false; - } - - const int64_t head_size = key->ne[0]; - const int64_t heads = key->ne[1]; - const int64_t n_tokens = key->ne[2]; - const int64_t n_embd = head_size * heads; - - if (head_size != 64 && head_size != 128) { - return false; - } - if (rk_weight->ne[0] != head_size || rk_weight->ne[1] != heads || ggml_nelements(rk_weight) != head_size * heads) { - return false; - } - if (sum_rows->ne[0] != 1 || sum_rows->ne[1] != heads || sum_rows->ne[2] != n_tokens) { - return false; - } - if (add->ne[0] != n_embd || add->ne[1] != n_tokens || ggml_nelements(add) != n_embd * n_tokens) { - return false; - } - if (!ggml_are_same_shape(add_cur, add) || ggml_nelements(reshape) != ggml_nelements(add)) { - return false; - } - - *cur = add_cur; - *k = key; - *r = rec; - *v = value; - *r_k = rk_weight; - return true; -} - -static int ggml_cuda_find_single_consumer(const ggml_cgraph * cgraph, int node_idx) { - if (ggml_node_get_use_count(cgraph, node_idx) != 1) { - return -1; - } - - const ggml_tensor * node = cgraph->nodes[node_idx]; - for (int i = node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * consumer = cgraph->nodes[i]; - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - if (consumer->src[src_idx] == node) { - return i; - } - } - } - - return -1; -} - static bool ggml_cuda_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b) { if (a == nullptr || b == nullptr || a->buffer == nullptr || b->buffer == nullptr) { return false; @@ -3844,23 +3675,6 @@ static bool ggml_cuda_add_mul_memory_ok( can_read_write_inplace(scale); } -static bool ggml_cuda_rwkv_rk_memory_ok( - const ggml_tensor * cur, - const ggml_tensor * k, - const ggml_tensor * r, - const ggml_tensor * v, - const ggml_tensor * r_k, - const ggml_tensor * dst) { - if (ggml_cuda_tensors_overlap(dst, cur) && !ggml_are_same_layout(dst, cur)) { - return false; - } - - return !ggml_cuda_tensors_overlap(dst, k) && - !ggml_cuda_tensors_overlap(dst, r) && - !ggml_cuda_tensors_overlap(dst, v) && - !ggml_cuda_tensors_overlap(dst, r_k); -} - // returns whether the write (out) nodes overwrite the read nodes in operation static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, @@ -4241,38 +4055,6 @@ static int ggml_cuda_try_fuse( return 2; } - // RWKV lerp: cur + (x_prev - cur) * weight. Time-mix has an extra REPEAT - // between SUB and MUL; channel-mix uses SUB directly. - if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_REPEAT, GGML_OP_MUL, GGML_OP_ADD }, { i + 3 })) { - const ggml_tensor * x_prev = nullptr; - const ggml_tensor * cur = nullptr; - const ggml_tensor * weight = nullptr; - - if (ggml_cuda_should_fuse_lerp(cgraph->nodes[i], cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], - &x_prev, &cur, &weight)) { - const int out_nodes[] = { i + 3 }; - if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 4, out_nodes, 1)) { - ggml_cuda_op_lerp_fused(*cuda_ctx, x_prev, cur, weight, cgraph->nodes[i + 3]); - return 3; - } - } - } - - if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_MUL, GGML_OP_ADD }, { i + 2 })) { - const ggml_tensor * x_prev = nullptr; - const ggml_tensor * cur = nullptr; - const ggml_tensor * weight = nullptr; - - if (ggml_cuda_should_fuse_lerp(cgraph->nodes[i], nullptr, cgraph->nodes[i + 1], cgraph->nodes[i + 2], - &x_prev, &cur, &weight)) { - const int out_nodes[] = { i + 2 }; - if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { - ggml_cuda_op_lerp_fused(*cuda_ctx, x_prev, cur, weight, cgraph->nodes[i + 2]); - return 2; - } - } - } - // RWKV key update: k + (a * ka - ka). if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_SUB, GGML_OP_ADD }, { i + 2 })) { const ggml_tensor * base = nullptr; @@ -4303,49 +4085,6 @@ static int ggml_cuda_try_fuse( } } - // RWKV r_k correction: - // cur + reshape(v * sum_rows((k * r) * r_k)) - if (node->op == GGML_OP_MUL) { - const int mul_kr_rk_idx = ggml_cuda_find_single_consumer(cgraph, i); - const int sum_rows_idx = mul_kr_rk_idx > i ? ggml_cuda_find_single_consumer(cgraph, mul_kr_rk_idx) : -1; - const int mul_v_rk_idx = sum_rows_idx > i ? ggml_cuda_find_single_consumer(cgraph, sum_rows_idx) : -1; - const int reshape_idx = mul_v_rk_idx > i ? ggml_cuda_find_single_consumer(cgraph, mul_v_rk_idx) : -1; - const int add_idx = reshape_idx > i ? ggml_cuda_find_single_consumer(cgraph, reshape_idx) : -1; - const ggml_tensor * cur = nullptr; - const ggml_tensor * k = nullptr; - const ggml_tensor * r = nullptr; - const ggml_tensor * v = nullptr; - const ggml_tensor * r_k = nullptr; - - const bool chain_ok = add_idx > i && - cgraph->nodes[mul_kr_rk_idx]->op == GGML_OP_MUL && - cgraph->nodes[sum_rows_idx]->op == GGML_OP_SUM_ROWS && - cgraph->nodes[mul_v_rk_idx]->op == GGML_OP_MUL && - cgraph->nodes[reshape_idx]->op == GGML_OP_RESHAPE && - cgraph->nodes[add_idx]->op == GGML_OP_ADD; - - bool range_ok = chain_ok; - for (int j = i + 1; range_ok && j < add_idx; ++j) { - if (j == mul_kr_rk_idx || j == sum_rows_idx || j == mul_v_rk_idx || j == reshape_idx) { - continue; - } - const ggml_tensor * n = cgraph->nodes[j]; - range_ok = ggml_is_empty(n) || n->op == GGML_OP_RESHAPE || n->op == GGML_OP_TRANSPOSE || - n->op == GGML_OP_VIEW || n->op == GGML_OP_PERMUTE || n->op == GGML_OP_NONE || - (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0; - } - - if (range_ok && - ggml_cuda_should_fuse_rwkv_rk(cgraph->nodes[i], cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], - cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], cgraph->nodes[add_idx], - &cur, &k, &r, &v, &r_k)) { - if (ggml_cuda_rwkv_rk_memory_ok(cur, k, r, v, r_k, cgraph->nodes[add_idx])) { - ggml_cuda_op_rwkv_rk_fused(*cuda_ctx, cur, k, r, v, r_k, cgraph->nodes[add_idx]); - return add_idx - i; - } - } - } - // Snake activation: y = x + sin(a*x)^2 * inv_b // Naive 5-op decomposition emitted by frontends: mul -> sin -> sqr -> mul -> add if (ggml_can_fuse_subgraph(cgraph, i, @@ -5859,6 +5598,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_ARANGE: case GGML_OP_TIMESTEP_EMBEDDING: case GGML_OP_LEAKY_RELU: + case GGML_OP_RWKV_LERP: + case GGML_OP_RWKV_RK: case GGML_OP_RWKV_WKV6: case GGML_OP_GATED_LINEAR_ATTN: case GGML_OP_RWKV_WKV7: diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 70fae6302e85..1453e55e5ee7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2176,8 +2176,6 @@ struct ggml_backend_vk_context { // If there's no fusion, bit 0 is still set. int fused_ops_write_mask {}; bool fused_add_mul {}; - bool fused_lerp {}; - bool fused_rwkv_rk {}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; @@ -11978,52 +11976,30 @@ static void ggml_vk_add_mul(ggml_backend_vk_context * ctx, vk_context& subctx, g ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, dst_buf, scale_buf }, pc, elements); } -static bool ggml_vk_should_fuse_lerp( - const ggml_backend_vk_context * ctx, - const ggml_tensor * sub, - const ggml_tensor * repeat, - const ggml_tensor * mul, - const ggml_tensor * add, - const ggml_tensor ** x_prev, - const ggml_tensor ** cur, - const ggml_tensor ** weight); - -static bool ggml_vk_should_fuse_rwkv_rk( - const ggml_backend_vk_context * ctx, - const ggml_tensor * mul_kr, - const ggml_tensor * mul_kr_rk, - const ggml_tensor * sum_rows, - const ggml_tensor * mul_v_rk, - const ggml_tensor * reshape, - const ggml_tensor * add, - const ggml_tensor ** cur, - const ggml_tensor ** k, - const ggml_tensor ** r, - const ggml_tensor ** v, - const ggml_tensor ** r_k); - -static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor * mul); - -static int ggml_vk_find_single_consumer(const ggml_cgraph * cgraph, int node_idx); - -static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { - const ggml_tensor * sub = cgraph->nodes[node_idx]; - const ggml_tensor * repeat = ctx->num_additional_fused_ops == 3 ? cgraph->nodes[node_idx + 1] : nullptr; - const ggml_tensor * mul = cgraph->nodes[node_idx + ctx->num_additional_fused_ops - 1]; - ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; - - const ggml_tensor * x_prev = nullptr; - const ggml_tensor * cur = nullptr; - const ggml_tensor * weight = nullptr; - GGML_ASSERT(ggml_vk_should_fuse_lerp(ctx, sub, repeat, mul, dst, &x_prev, &cur, &weight)); +static void ggml_vk_rwkv_lerp(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * x_prev = dst->src[0]; + const ggml_tensor * cur = dst->src[1]; + const ggml_tensor * weight = dst->src[2]; + + GGML_ASSERT(x_prev->type == GGML_TYPE_F32 && cur->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(x_prev, cur)); + GGML_ASSERT(ggml_is_contiguous(x_prev) && ggml_is_contiguous(cur) && ggml_is_contiguous(weight) && ggml_is_contiguous(dst)); const uint32_t x_offset = get_misalign_bytes(ctx, x_prev) / ggml_type_size(x_prev->type); const uint32_t c_offset = get_misalign_bytes(ctx, cur) / ggml_type_size(cur->type); const uint32_t w_offset = get_misalign_bytes(ctx, weight) / ggml_type_size(weight->type); const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); - const uint32_t base_total = (uint32_t) ggml_nelements(x_prev); - const uint32_t n_mix = (uint32_t) (ggml_nelements(dst) / ggml_nelements(x_prev)); + const int64_t base_total_i64 = ggml_nelements(x_prev); + GGML_ASSERT(base_total_i64 > 0 && ggml_nelements(dst) % base_total_i64 == 0); + const int64_t n_mix_i64 = ggml_nelements(dst) / base_total_i64; + + GGML_ASSERT(dst->ne[0] == x_prev->ne[0] && dst->ne[1] == x_prev->ne[1] && dst->ne[2] == x_prev->ne[2]); + GGML_ASSERT(weight->ne[0] == dst->ne[0] && weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == n_mix_i64); + GGML_ASSERT(dst->ne[0] <= UINT32_MAX && base_total_i64 <= UINT32_MAX && n_mix_i64 <= UINT32_MAX); + + const uint32_t base_total = (uint32_t) base_total_i64; + const uint32_t n_mix = (uint32_t) n_mix_i64; vk_op_lerp_push_constants pc { (uint32_t) dst->ne[0], @@ -12044,25 +12020,18 @@ static void ggml_vk_lerp_fused(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_prev_buf, cur_buf, weight_buf, dst_buf }, pc, { base_total, n_mix, 1 }); } -static void ggml_vk_rwkv_rk_fused(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { - const ggml_tensor * mul_kr = cgraph->nodes[node_idx]; - const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); - const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); - const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); - const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); - GGML_ASSERT(mul_kr_rk_idx > node_idx && sum_rows_idx > node_idx && mul_v_rk_idx > node_idx && reshape_idx > node_idx); - const ggml_tensor * mul_kr_rk = cgraph->nodes[mul_kr_rk_idx]; - const ggml_tensor * sum_rows = cgraph->nodes[sum_rows_idx]; - const ggml_tensor * mul_v_rk = cgraph->nodes[mul_v_rk_idx]; - const ggml_tensor * reshape = cgraph->nodes[reshape_idx]; - ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; - - const ggml_tensor * cur = nullptr; - const ggml_tensor * k = nullptr; - const ggml_tensor * r = nullptr; - const ggml_tensor * v = nullptr; - const ggml_tensor * r_k = nullptr; - GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, mul_kr, mul_kr_rk, sum_rows, mul_v_rk, reshape, dst, &cur, &k, &r, &v, &r_k)); +static void ggml_vk_rwkv_rk(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * cur = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * r = dst->src[2]; + const ggml_tensor * v = dst->src[3]; + const ggml_tensor * r_k = dst->src[4]; + + GGML_ASSERT(cur->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32 && r->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32 && r_k->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(cur) && ggml_is_contiguous(k) && ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(v) && ggml_is_contiguous(r_k) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(k, r) && ggml_are_same_shape(k, v) && ggml_are_same_shape(cur, dst)); const uint32_t cur_offset = get_misalign_bytes(ctx, cur) / ggml_type_size(cur->type); const uint32_t k_offset = get_misalign_bytes(ctx, k) / ggml_type_size(k->type); @@ -12076,6 +12045,10 @@ static void ggml_vk_rwkv_rk_fused(ggml_backend_vk_context * ctx, vk_context& sub const uint32_t T = (uint32_t) k->ne[2]; const uint32_t C = head_size * H; + GGML_ASSERT(head_size == 64 || head_size == 128); + GGML_ASSERT(r_k->ne[0] == head_size && r_k->ne[1] == H && ggml_nelements(r_k) == head_size * H); + GGML_ASSERT(dst->ne[0] == C && dst->ne[1] == T && ggml_nelements(dst) == C * T); + vk_op_rwkv_rk_push_constants pc { C, H, @@ -14842,54 +14815,24 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr return false; }; - if (ctx->fused_rwkv_rk) { - const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); - const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); - const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); - const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); - const int add_idx = node_idx + ctx->num_additional_fused_ops; - - const ggml_tensor * cur = nullptr; - const ggml_tensor * k = nullptr; - const ggml_tensor * r = nullptr; - const ggml_tensor * v = nullptr; - const ggml_tensor * r_k = nullptr; - GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], - cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], - cgraph->nodes[add_idx], &cur, &k, &r, &v, &r_k)); - - const ggml_tensor * dst = cgraph->nodes[add_idx]; - const ggml_tensor * fused_srcs[] = { cur, k, r, v, r_k }; - - if (overlaps_unsynced(dst, ctx->unsynced_nodes_read) || overlaps_unsynced(dst, ctx->unsynced_nodes_written)) { - need_sync = true; - } - for (const ggml_tensor * src : fused_srcs) { - if (src && overlaps_unsynced(src, ctx->unsynced_nodes_written)) { + // For all fused ops, check if the destination node or any of the source + // nodes require synchronization. + for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1 && !need_sync; ++i) { + const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; + // If the node actually writes to memory, then check if it needs to sync + if (ctx->fused_ops_write_mask & (1 << i)) { + if (overlaps_unsynced(cur_node, ctx->unsynced_nodes_read) || overlaps_unsynced(cur_node, ctx->unsynced_nodes_written)) { need_sync = true; break; } } - } else { - // For all fused ops, check if the destination node or any of the source - // nodes require synchronization. - for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1 && !need_sync; ++i) { - const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; - // If the node actually writes to memory, then check if it needs to sync - if (ctx->fused_ops_write_mask & (1 << i)) { - if (overlaps_unsynced(cur_node, ctx->unsynced_nodes_read) || overlaps_unsynced(cur_node, ctx->unsynced_nodes_written)) { - need_sync = true; - break; - } + for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { + if (!cur_node->src[j]) { + continue; } - for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { - if (!cur_node->src[j]) { - continue; - } - if (overlaps_unsynced(cur_node->src[j], ctx->unsynced_nodes_written)) { - need_sync = true; - break; - } + if (overlaps_unsynced(cur_node->src[j], ctx->unsynced_nodes_written)) { + need_sync = true; + break; } } } @@ -14908,42 +14851,18 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr ggml_vk_sync_buffers(ctx, compute_ctx); } } - if (ctx->fused_rwkv_rk) { - const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); - const int sum_rows_idx = ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx); - const int mul_v_rk_idx = ggml_vk_find_single_consumer(cgraph, sum_rows_idx); - const int reshape_idx = ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx); - const int add_idx = node_idx + ctx->num_additional_fused_ops; - - const ggml_tensor * cur = nullptr; - const ggml_tensor * k = nullptr; - const ggml_tensor * r = nullptr; - const ggml_tensor * v = nullptr; - const ggml_tensor * r_k = nullptr; - GGML_ASSERT(ggml_vk_should_fuse_rwkv_rk(ctx, node, cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], - cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], - cgraph->nodes[add_idx], &cur, &k, &r, &v, &r_k)); - - ctx->unsynced_nodes_written.push_back(cgraph->nodes[add_idx]); - ctx->unsynced_nodes_read.push_back(cur); - ctx->unsynced_nodes_read.push_back(k); - ctx->unsynced_nodes_read.push_back(r); - ctx->unsynced_nodes_read.push_back(v); - ctx->unsynced_nodes_read.push_back(r_k); - } else { - // Add all fused nodes to the unsynchronized lists. - for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { - const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; - // Multiple outputs could be written, e.g. in topk_moe. Add them all to the list. - if (ctx->fused_ops_write_mask & (1 << i)) { - ctx->unsynced_nodes_written.push_back(cur_node); - } - for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { - if (!cur_node->src[j]) { - continue; - } - ctx->unsynced_nodes_read.push_back(cur_node->src[j]); + // Add all fused nodes to the unsynchronized lists. + for (int32_t i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { + const ggml_tensor *cur_node = cgraph->nodes[node_idx + i]; + // Multiple outputs could be written, e.g. in topk_moe. Add them all to the list. + if (ctx->fused_ops_write_mask & (1 << i)) { + ctx->unsynced_nodes_written.push_back(cur_node); + } + for (uint32_t j = 0; j < GGML_MAX_SRC; ++j) { + if (!cur_node->src[j]) { + continue; } + ctx->unsynced_nodes_read.push_back(cur_node->src[j]); } } } @@ -14994,17 +14913,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr } break; case GGML_OP_SUB: - if (ctx->fused_lerp) { - ggml_vk_lerp_fused(ctx, compute_ctx, cgraph, node_idx); - } else { - ggml_vk_sub(ctx, compute_ctx, src0, src1, node); - } + ggml_vk_sub(ctx, compute_ctx, src0, src1, node); break; case GGML_OP_MUL: - if (ctx->fused_rwkv_rk) { - ggml_vk_rwkv_rk_fused(ctx, compute_ctx, cgraph, node_idx); - } else if (ctx->num_additional_fused_ops) { + if (ctx->num_additional_fused_ops) { ggml_vk_snake_dispatch_fused(ctx, compute_ctx, cgraph, node_idx); } else { ggml_vk_mul(ctx, compute_ctx, src0, src1, node); @@ -15290,6 +15203,16 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_RWKV_LERP: + ggml_vk_rwkv_lerp(ctx, compute_ctx, node); + + break; + + case GGML_OP_RWKV_RK: + ggml_vk_rwkv_rk(ctx, compute_ctx, node); + + break; + case GGML_OP_RWKV_WKV6: ggml_vk_rwkv_wkv6(ctx, compute_ctx, node); @@ -16086,203 +16009,39 @@ static bool ggml_vk_should_fuse_add_mul(const ggml_backend_vk_context * ctx, con return true; } -static bool ggml_vk_should_fuse_lerp( - const ggml_backend_vk_context * ctx, - const ggml_tensor * sub, - const ggml_tensor * repeat, - const ggml_tensor * mul, - const ggml_tensor * add, - const ggml_tensor ** x_prev, - const ggml_tensor ** cur, - const ggml_tensor ** weight) { - if (sub->op != GGML_OP_SUB || mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD) { - return false; - } +// Check whether the tensors overlap in memory. +// Fusions can potentially overwrite src tensors in ways that are not prevented +// by ggml-alloc. If the fusion src is being applied in a way that's elementwise +// with the destination, then it's OK for them to overlap if they are exactly equal. +static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b, bool elementwise) { + ggml_backend_vk_buffer_context * a_buf_ctx = (ggml_backend_vk_buffer_context *)a->buffer->context; + vk_buffer a_buf = a_buf_ctx->dev_buffer; + ggml_backend_vk_buffer_context * b_buf_ctx = (ggml_backend_vk_buffer_context *)b->buffer->context; + vk_buffer b_buf = b_buf_ctx->dev_buffer; + if (a_buf == b_buf) { + auto a_base = vk_tensor_offset(a) + a->view_offs; + auto a_size = ggml_nbytes(a); + auto b_base = vk_tensor_offset(b) + b->view_offs; + auto b_size = ggml_nbytes(b); - const ggml_tensor * sx = sub; - if (repeat) { - if (repeat->op != GGML_OP_REPEAT || repeat->src[0] != sub) { + if (elementwise && a_base == b_base && a_size == b_size) { return false; } - sx = repeat; - } - const ggml_tensor * w = nullptr; - if (mul->src[0] == sx) { - w = mul->src[1]; - } else if (mul->src[1] == sx) { - w = mul->src[0]; - } else { - return false; - } - - const ggml_tensor * add_cur = nullptr; - if (add->src[0] == mul) { - add_cur = add->src[1]; - } else if (add->src[1] == mul) { - add_cur = add->src[0]; - } else { - return false; - } - - if (add_cur != sub->src[1]) { - return false; - } - - const ggml_tensor * xp = sub->src[0]; - const ggml_tensor * c = sub->src[1]; - if (xp->type != GGML_TYPE_F32 || c->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { - return false; - } - - if (!ggml_can_repeat(xp, add) || !ggml_can_repeat(c, add) || !ggml_can_repeat(w, add)) { - return false; - } - - if (!ggml_is_contiguous(xp) || !ggml_is_contiguous(c) || !ggml_is_contiguous(w) || !ggml_is_contiguous(add)) { - return false; - } - - const int64_t base_total = ggml_nelements(xp); - const int64_t total = ggml_nelements(add); - if (!ggml_are_same_shape(xp, c) || - base_total <= 0 || - total % base_total != 0 || - add->ne[0] != xp->ne[0] || - add->ne[1] != xp->ne[1] || - add->ne[2] != xp->ne[2]) { - return false; - } - - const int64_t n_mix = total / base_total; - if (w->ne[0] != add->ne[0] || - w->ne[1] != 1 || - w->ne[2] != 1 || - w->ne[3] != n_mix || - ggml_nelements(w) != add->ne[0] * n_mix) { - return false; - } - - if (add->ne[0] > UINT32_MAX || base_total > UINT32_MAX || n_mix > UINT32_MAX) { - return false; - } - - if (get_misalign_bytes(ctx, xp) / ggml_type_size(xp->type) > UINT32_MAX || - get_misalign_bytes(ctx, c) / ggml_type_size(c->type) > UINT32_MAX || - get_misalign_bytes(ctx, w) / ggml_type_size(w->type) > UINT32_MAX || - get_misalign_bytes(ctx, add) / ggml_type_size(add->type) > UINT32_MAX) { - return false; + if ((b_base <= a_base && a_base < b_base + b_size) || + (a_base <= b_base && b_base < a_base + a_size)) { + return true; + } } - - *x_prev = xp; - *cur = c; - *weight = w; - return true; + return false; } -static bool ggml_vk_should_fuse_rwkv_rk( - const ggml_backend_vk_context * ctx, - const ggml_tensor * mul_kr, - const ggml_tensor * mul_kr_rk, - const ggml_tensor * sum_rows, - const ggml_tensor * mul_v_rk, - const ggml_tensor * reshape, - const ggml_tensor * add, - const ggml_tensor ** cur, - const ggml_tensor ** k, - const ggml_tensor ** r, - const ggml_tensor ** v, - const ggml_tensor ** r_k) { - if (mul_kr->op != GGML_OP_MUL || mul_kr_rk->op != GGML_OP_MUL || - sum_rows->op != GGML_OP_SUM_ROWS || mul_v_rk->op != GGML_OP_MUL || - reshape->op != GGML_OP_RESHAPE || add->op != GGML_OP_ADD) { - return false; - } - - if (sum_rows->src[0] != mul_kr_rk || reshape->src[0] != mul_v_rk) { - return false; - } - - const ggml_tensor * rk_weight = nullptr; - if (mul_kr_rk->src[0] == mul_kr) { - rk_weight = mul_kr_rk->src[1]; - } else if (mul_kr_rk->src[1] == mul_kr) { - rk_weight = mul_kr_rk->src[0]; - } else { - return false; - } - - const ggml_tensor * value = nullptr; - if (mul_v_rk->src[0] == sum_rows) { - value = mul_v_rk->src[1]; - } else if (mul_v_rk->src[1] == sum_rows) { - value = mul_v_rk->src[0]; - } else { - return false; - } - - const ggml_tensor * add_cur = nullptr; - if (add->src[0] == reshape) { - add_cur = add->src[1]; - } else if (add->src[1] == reshape) { - add_cur = add->src[0]; - } else { - return false; - } - - const ggml_tensor * key = mul_kr->src[0]; - const ggml_tensor * rec = mul_kr->src[1]; - if (key->type != GGML_TYPE_F32 || rec->type != GGML_TYPE_F32 || value->type != GGML_TYPE_F32 || - rk_weight->type != GGML_TYPE_F32 || add_cur->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32) { - return false; - } - - if (!ggml_is_contiguous(key) || !ggml_is_contiguous(rec) || !ggml_is_contiguous(value) || - !ggml_is_contiguous(rk_weight) || !ggml_is_contiguous(add_cur) || !ggml_is_contiguous(add)) { - return false; - } - - if (!ggml_are_same_shape(key, rec) || !ggml_are_same_shape(key, value)) { - return false; - } - - const int64_t head_size = key->ne[0]; - const int64_t heads = key->ne[1]; - const int64_t n_tokens = key->ne[2]; - const int64_t n_embd = head_size * heads; - - if (head_size != 64 && head_size != 128) { - return false; - } - if (rk_weight->ne[0] != head_size || rk_weight->ne[1] != heads || ggml_nelements(rk_weight) != head_size * heads) { - return false; - } - if (sum_rows->ne[0] != 1 || sum_rows->ne[1] != heads || sum_rows->ne[2] != n_tokens) { - return false; - } - if (add->ne[0] != n_embd || add->ne[1] != n_tokens || ggml_nelements(add) != n_embd * n_tokens) { - return false; - } - if (!ggml_are_same_shape(add_cur, add) || ggml_nelements(reshape) != ggml_nelements(add)) { - return false; - } - if (n_embd > UINT32_MAX || heads > UINT32_MAX || n_tokens > UINT32_MAX) { - return false; - } - - const ggml_tensor * tensors[] = { add_cur, key, rec, value, rk_weight, add }; - for (const ggml_tensor * tensor : tensors) { - if (get_misalign_bytes(ctx, tensor) / ggml_type_size(tensor->type) > UINT32_MAX) { - return false; - } - } +static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor * mul) { + const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; - *cur = add_cur; - *k = key; - *r = rec; - *v = value; - *r_k = rk_weight; - return true; + return !ggml_vk_tensors_overlap(add->src[0], mul, true) && + !ggml_vk_tensors_overlap(add->src[1], mul, true) && + !ggml_vk_tensors_overlap(scale, mul, true); } static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { @@ -16722,117 +16481,6 @@ static bool ggml_vk_can_fuse_snake(ggml_backend_vk_context * ctx, const struct g return true; } -// Check whether the tensors overlap in memory. -// Fusions can potentially overwrite src tensors in ways that are not prevented -// by ggml-alloc. If the fusion src is being applied in a way that's elementwise -// with the destination, then it's OK for them to overlap if they are exactly equal. -static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b, bool elementwise) { - ggml_backend_vk_buffer_context * a_buf_ctx = (ggml_backend_vk_buffer_context *)a->buffer->context; - vk_buffer a_buf = a_buf_ctx->dev_buffer; - ggml_backend_vk_buffer_context * b_buf_ctx = (ggml_backend_vk_buffer_context *)b->buffer->context; - vk_buffer b_buf = b_buf_ctx->dev_buffer; - if (a_buf == b_buf) { - auto a_base = vk_tensor_offset(a) + a->view_offs; - auto a_size = ggml_nbytes(a); - auto b_base = vk_tensor_offset(b) + b->view_offs; - auto b_size = ggml_nbytes(b); - - if (elementwise && a_base == b_base && a_size == b_size) { - return false; - } - - if ((b_base <= a_base && a_base < b_base + b_size) || - (a_base <= b_base && b_base < a_base + a_size)) { - return true; - } - } - return false; -} - -static int ggml_vk_find_single_consumer(const ggml_cgraph * cgraph, int node_idx) { - if (ggml_node_get_use_count(cgraph, node_idx) != 1) { - return -1; - } - - const ggml_tensor * node = cgraph->nodes[node_idx]; - for (int i = node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * consumer = cgraph->nodes[i]; - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - if (consumer->src[src_idx] == node) { - return i; - } - } - } - - return -1; -} - -static bool ggml_vk_add_mul_memory_ok(const ggml_tensor * add, const ggml_tensor * mul) { - const ggml_tensor * scale = mul->src[0] == add ? mul->src[1] : mul->src[0]; - - return !ggml_vk_tensors_overlap(add->src[0], mul, true) && - !ggml_vk_tensors_overlap(add->src[1], mul, true) && - !ggml_vk_tensors_overlap(scale, mul, true); -} - -static bool ggml_vk_can_fuse_rwkv_rk( - const ggml_backend_vk_context * ctx, - const ggml_cgraph * cgraph, - int node_idx, - int * add_idx) { - const int mul_kr_rk_idx = ggml_vk_find_single_consumer(cgraph, node_idx); - const int sum_rows_idx = mul_kr_rk_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, mul_kr_rk_idx) : -1; - const int mul_v_rk_idx = sum_rows_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, sum_rows_idx) : -1; - const int reshape_idx = mul_v_rk_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, mul_v_rk_idx) : -1; - const int found_add_idx = reshape_idx > node_idx ? ggml_vk_find_single_consumer(cgraph, reshape_idx) : -1; - - const bool chain_ok = found_add_idx > node_idx && - cgraph->nodes[node_idx]->op == GGML_OP_MUL && - cgraph->nodes[mul_kr_rk_idx]->op == GGML_OP_MUL && - cgraph->nodes[sum_rows_idx]->op == GGML_OP_SUM_ROWS && - cgraph->nodes[mul_v_rk_idx]->op == GGML_OP_MUL && - cgraph->nodes[reshape_idx]->op == GGML_OP_RESHAPE && - cgraph->nodes[found_add_idx]->op == GGML_OP_ADD; - if (!chain_ok) { - return false; - } - - for (int i = node_idx + 1; i < found_add_idx; ++i) { - if (i == mul_kr_rk_idx || i == sum_rows_idx || i == mul_v_rk_idx || i == reshape_idx) { - continue; - } - const ggml_tensor * node = cgraph->nodes[i]; - if (!ggml_is_empty(node) && node->op != GGML_OP_RESHAPE && node->op != GGML_OP_TRANSPOSE && - node->op != GGML_OP_VIEW && node->op != GGML_OP_PERMUTE && node->op != GGML_OP_NONE && - (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0) { - return false; - } - } - - const ggml_tensor * cur = nullptr; - const ggml_tensor * k = nullptr; - const ggml_tensor * r = nullptr; - const ggml_tensor * v = nullptr; - const ggml_tensor * r_k = nullptr; - if (!ggml_vk_should_fuse_rwkv_rk(ctx, cgraph->nodes[node_idx], cgraph->nodes[mul_kr_rk_idx], cgraph->nodes[sum_rows_idx], - cgraph->nodes[mul_v_rk_idx], cgraph->nodes[reshape_idx], cgraph->nodes[found_add_idx], - &cur, &k, &r, &v, &r_k)) { - return false; - } - - const ggml_tensor * dst = cgraph->nodes[found_add_idx]; - if (ggml_vk_tensors_overlap(cur, dst, true) || - ggml_vk_tensors_overlap(k, dst, true) || - ggml_vk_tensors_overlap(r, dst, true) || - ggml_vk_tensors_overlap(v, dst, true) || - ggml_vk_tensors_overlap(r_k, dst, true)) { - return false; - } - - *add_idx = found_add_idx; - return true; -} - static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { GGML_UNUSED(ctx); @@ -17033,10 +16681,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; - ctx->fused_lerp = false; - ctx->fused_rwkv_rk = false; const char *fusion_string {}; - int rwkv_rk_add_idx = -1; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); if (num_adds) { @@ -17070,39 +16715,12 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_REPEAT, GGML_OP_MUL, GGML_OP_ADD }, { i + 3 })) { - const ggml_tensor * x_prev = nullptr; - const ggml_tensor * cur = nullptr; - const ggml_tensor * weight = nullptr; - if (ggml_vk_should_fuse_lerp(ctx, cgraph->nodes[i], cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], - &x_prev, &cur, &weight)) { - ctx->num_additional_fused_ops = 3; - ctx->fused_lerp = true; - fusion_string = "LERP"; - std::fill_n(op_srcs_fused_elementwise, 4, true); - } - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SUB, GGML_OP_MUL, GGML_OP_ADD }, { i + 2 })) { - const ggml_tensor * x_prev = nullptr; - const ggml_tensor * cur = nullptr; - const ggml_tensor * weight = nullptr; - if (ggml_vk_should_fuse_lerp(ctx, cgraph->nodes[i], nullptr, cgraph->nodes[i + 1], cgraph->nodes[i + 2], - &x_prev, &cur, &weight)) { - ctx->num_additional_fused_ops = 2; - ctx->fused_lerp = true; - fusion_string = "LERP"; - std::fill_n(op_srcs_fused_elementwise, 3, true); - } } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_ADD, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; ctx->fused_add_mul = true; fusion_string = "ADD_MUL"; op_srcs_fused_elementwise[0] = true; op_srcs_fused_elementwise[1] = true; - } else if (cgraph->nodes[i]->op == GGML_OP_MUL && ggml_vk_can_fuse_rwkv_rk(ctx, cgraph, i, &rwkv_rk_add_idx)) { - ctx->fused_rwkv_rk = true; - ctx->num_additional_fused_ops = rwkv_rk_add_idx - i; - fusion_string = "RWKV_RK"; - std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false); } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD })) { ctx->num_additional_fused_ops = 2; fusion_string = "NORM_MUL_ADD"; @@ -17272,8 +16890,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_add_mul = false; - ctx->fused_lerp = false; - ctx->fused_rwkv_rk = false; } } @@ -17324,8 +16940,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 0; ctx->fused_add_mul = false; - ctx->fused_lerp = false; - ctx->fused_rwkv_rk = false; } ctx->last_total_flops = total_flops; @@ -18372,6 +17986,86 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm && op->src[1]->type == GGML_TYPE_F32; case GGML_OP_POOL_2D: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; + case GGML_OP_RWKV_LERP: + { + const ggml_tensor * x_prev = op->src[0]; + const ggml_tensor * cur = op->src[1]; + const ggml_tensor * weight = op->src[2]; + + if (x_prev == nullptr || cur == nullptr || weight == nullptr) { + return false; + } + if (x_prev->type != GGML_TYPE_F32 || cur->type != GGML_TYPE_F32 || weight->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_are_same_shape(x_prev, cur) || + !ggml_can_repeat(x_prev, op) || + !ggml_can_repeat(cur, op) || + !ggml_can_repeat(weight, op)) { + return false; + } + if (!ggml_is_contiguous(x_prev) || !ggml_is_contiguous(cur) || !ggml_is_contiguous(weight) || !ggml_is_contiguous(op)) { + return false; + } + + const int64_t base_total = ggml_nelements(x_prev); + const int64_t total = ggml_nelements(op); + if (base_total <= 0 || total % base_total != 0 || + op->ne[0] != x_prev->ne[0] || + op->ne[1] != x_prev->ne[1] || + op->ne[2] != x_prev->ne[2]) { + return false; + } + + const int64_t n_mix = total / base_total; + if (weight->ne[0] != op->ne[0] || + weight->ne[1] != 1 || + weight->ne[2] != 1 || + weight->ne[3] != n_mix || + ggml_nelements(weight) != op->ne[0] * n_mix) { + return false; + } + + return op->ne[0] <= UINT32_MAX && base_total <= UINT32_MAX && n_mix <= UINT32_MAX; + } + case GGML_OP_RWKV_RK: + { + const ggml_tensor * cur = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * r = op->src[2]; + const ggml_tensor * v = op->src[3]; + const ggml_tensor * r_k = op->src[4]; + + const ggml_tensor * tensors[] = { cur, k, r, v, r_k, op }; + for (const ggml_tensor * tensor : tensors) { + if (tensor == nullptr || tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor)) { + return false; + } + } + if (!ggml_are_same_shape(k, r) || !ggml_are_same_shape(k, v)) { + return false; + } + + const int64_t head_size = k->ne[0]; + const int64_t heads = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_embd = head_size * heads; + + if (head_size != 64 && head_size != 128) { + return false; + } + if (r_k->ne[0] != head_size || r_k->ne[1] != heads || ggml_nelements(r_k) != head_size * heads) { + return false; + } + if (op->ne[0] != n_embd || op->ne[1] != n_tokens || ggml_nelements(op) != n_embd * n_tokens) { + return false; + } + if (!ggml_are_same_shape(cur, op)) { + return false; + } + + return n_embd <= UINT32_MAX && heads <= UINT32_MAX && n_tokens <= UINT32_MAX; + } case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: return true; // all inputs are contiguous, see ggml.c @@ -19328,6 +19022,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * } else if (tensor->op == GGML_OP_LEAKY_RELU) { const float * op_params = (const float *)tensor->op_params; tensor_clone = ggml_leaky_relu(ggml_ctx, src_clone[0], op_params[0], false); + } else if (tensor->op == GGML_OP_RWKV_LERP) { + tensor_clone = ggml_rwkv_lerp(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]); + } else if (tensor->op == GGML_OP_RWKV_RK) { + tensor_clone = ggml_rwkv_rk(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4]); } else if (tensor->op == GGML_OP_RWKV_WKV6) { tensor_clone = ggml_rwkv_wkv6(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5]); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 5425a7b5878d..2ee5c72ff910 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1056,6 +1056,8 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "WIN_UNPART", "GET_REL_POS", "ADD_REL_POS", + "RWKV_LERP", + "RWKV_RK", "RWKV_WKV6", "GATED_LINEAR_ATTN", "RWKV_WKV7", @@ -1078,7 +1080,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1167,6 +1169,8 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "win_unpart(x)", "get_rel_pos(x)", "add_rel_pos(x)", + "rwkv_lerp(x_prev, cur, weight)", + "rwkv_rk(cur, k, r, v, r_k)", "rwkv_wkv6(k, v, r, tf, td, s)", "gated_linear_attn(k, v, q, gate, s)", "rwkv_wkv7(r, w, k, v, kk, a, s)", @@ -1189,7 +1193,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -5718,6 +5722,91 @@ struct ggml_tensor * ggml_add_rel_pos_inplace( return ggml_add_rel_pos_impl(ctx, a, pw, ph, true); } +// ggml_rwkv_lerp + +struct ggml_tensor * ggml_rwkv_lerp( + struct ggml_context * ctx, + struct ggml_tensor * x_prev, + struct ggml_tensor * cur, + struct ggml_tensor * weight) { + GGML_ASSERT(x_prev->type == GGML_TYPE_F32); + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(weight->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_are_same_shape(x_prev, cur)); + GGML_ASSERT(ggml_is_contiguous(x_prev)); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(weight)); + + const int64_t ne[4] = { + cur->ne[0], + cur->ne[1], + cur->ne[2], + MAX(cur->ne[3], weight->ne[3]), + }; + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + GGML_ASSERT(ggml_can_repeat(x_prev, result)); + GGML_ASSERT(ggml_can_repeat(cur, result)); + GGML_ASSERT(ggml_can_repeat(weight, result)); + + result->op = GGML_OP_RWKV_LERP; + result->src[0] = x_prev; + result->src[1] = cur; + result->src[2] = weight; + + return result; +} + +// ggml_rwkv_rk + +struct ggml_tensor * ggml_rwkv_rk( + struct ggml_context * ctx, + struct ggml_tensor * cur, + struct ggml_tensor * k, + struct ggml_tensor * r, + struct ggml_tensor * v, + struct ggml_tensor * r_k) { + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(r->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(r_k->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r_k)); + + GGML_ASSERT(ggml_are_same_shape(k, r)); + GGML_ASSERT(ggml_are_same_shape(k, v)); + + const int64_t head_size = k->ne[0]; + const int64_t n_heads = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_embd = head_size * n_heads; + + GGML_ASSERT(r_k->ne[0] == head_size); + GGML_ASSERT(r_k->ne[1] == n_heads); + GGML_ASSERT(ggml_nelements(r_k) == head_size * n_heads); + GGML_ASSERT(cur->ne[0] == n_embd); + GGML_ASSERT(cur->ne[1] == n_tokens); + GGML_ASSERT(ggml_nelements(cur) == n_embd * n_tokens); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, cur->ne); + + result->op = GGML_OP_RWKV_RK; + result->src[0] = cur; + result->src[1] = k; + result->src[2] = r; + result->src[3] = v; + result->src[4] = r_k; + + return result; +} + // ggml_rwkv_wkv6 struct ggml_tensor * ggml_rwkv_wkv6( diff --git a/src/models/rwkv7-base.cpp b/src/models/rwkv7-base.cpp index 50b24e5dccfe..4bb44925ba2a 100644 --- a/src/models/rwkv7-base.cpp +++ b/src/models/rwkv7-base.cpp @@ -48,11 +48,7 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in bool has_gating = layer.time_mix_g1 && layer.time_mix_g2; - ggml_tensor * sx = ggml_sub(ctx0, x_prev, cur); - ggml_tensor * dummy = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_embd, n_seq_tokens, n_seqs, has_gating ? 6 : 5); - sx = ggml_repeat(ctx0, sx, dummy); - - ggml_tensor * xxx = ggml_add(ctx0, ggml_mul(ctx0, sx, layer.time_mix_lerp_fused), cur); + ggml_tensor * xxx = ggml_rwkv_lerp(ctx0, x_prev, cur, layer.time_mix_lerp_fused); ggml_tensor * xr = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], 0); ggml_tensor * xw = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float)); @@ -124,8 +120,7 @@ ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * in } else { cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens); } - ggml_tensor * rk = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_mul(ctx0, k, r), r_k)); - cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, ggml_mul(ctx0, v, rk), n_embd, n_tokens)); + cur = ggml_rwkv_rk(ctx0, cur, k, r, v, r_k); if (has_gating) { cur = ggml_mul(ctx0, cur, g); From bdbd216d362819f8b0e0d4e68558986b09e6d854 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 21:20:22 +0800 Subject: [PATCH 12/16] vulkan: disable rwkv7 t1 shader on Intel Windows --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 1453e55e5ee7..6e608a444749 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5453,7 +5453,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1); - if (device->subgroup_arithmetic && device->subgroup_size <= 64 && 64 % device->subgroup_size == 0) { + // Intel Windows fails RWKV_WKV7 T=1 correctness with this subgroup shader. + if (device->subgroup_arithmetic && + device->driver_id != vk::DriverId::eIntelProprietaryWindows && + device->subgroup_size <= 64 && 64 % device->subgroup_size == 0) { const uint32_t rows_per_wg = 4; ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_t1_f32, "rwkv_wkv7_t1_f32", rwkv_wkv7_t1_f32_len, rwkv_wkv7_t1_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, rows_per_wg, 1}, {device->subgroup_size, rows_per_wg}, 1, true, true, device->subgroup_size); } From 03d17d573a05f510bbbe8ae4e58c66224a3b35c7 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 21:52:11 +0800 Subject: [PATCH 13/16] tests: add rwkv lerp and rk backend ops --- tests/test-backend-ops.cpp | 59 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4cb160cb2e7e..74eb70da1d86 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3850,6 +3850,27 @@ struct test_ssm_scan : public test_case { } }; +// GGML_OP_RWKV_LERP +struct test_rwkv_lerp : public test_case { + const int64_t n_embd; + const int64_t n_seq_tokens; + const int64_t n_mix; + + std::string vars() override { + return VARS_TO_STR3(n_embd, n_seq_tokens, n_mix); + } + + test_rwkv_lerp(int64_t n_embd = 2048, int64_t n_seq_tokens = 32, int64_t n_mix = 6) + : n_embd(n_embd), n_seq_tokens(n_seq_tokens), n_mix(n_mix) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * x_prev = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_seq_tokens); + ggml_tensor * cur = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_seq_tokens); + ggml_tensor * weight = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd, 1, 1, n_mix); + return ggml_rwkv_lerp(ctx, x_prev, cur, weight); + } +}; + // GGML_OP_RWKV_WKV6 struct test_rwkv_wkv6 : public test_case { const ggml_type type; @@ -3880,6 +3901,36 @@ struct test_rwkv_wkv6 : public test_case { } }; +// GGML_OP_RWKV_RK +struct test_rwkv_rk : public test_case { + const int64_t head_count; + const int64_t head_size; + const int64_t n_seq_tokens; + + std::string vars() override { + return VARS_TO_STR3(head_count, head_size, n_seq_tokens); + } + + test_rwkv_rk(int64_t head_count = 32, int64_t head_size = 64, int64_t n_seq_tokens = 32) + : head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t n_embd = head_size * head_count; + ggml_tensor * cur = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_seq_tokens); + ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, n_seq_tokens); + ggml_tensor * r = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, n_seq_tokens); + ggml_tensor * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, n_seq_tokens); + ggml_tensor * r_k = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, head_size, head_count); + return ggml_rwkv_rk(ctx, cur, k, r, v, r_k); + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + init_tensor_uniform(t, -0.25f, 0.25f); + } + } +}; + // GGML_OP_GATED_DELTA_NET struct test_gated_delta_net : public test_case { const ggml_type type; @@ -8405,6 +8456,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap + test_cases.emplace_back(new test_rwkv_lerp(2048, 1, 6)); + test_cases.emplace_back(new test_rwkv_lerp(2048, 32, 6)); + + test_cases.emplace_back(new test_rwkv_rk(32, 64, 1)); + test_cases.emplace_back(new test_rwkv_rk(32, 64, 32)); + test_cases.emplace_back(new test_rwkv_rk(16, 128, 1)); + test_cases.emplace_back(new test_rwkv_rk(16, 128, 32)); + test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 4)); From 7839033195201252c75ddfbb9bb689928b893479 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Thu, 2 Jul 2026 22:06:16 +0800 Subject: [PATCH 14/16] cuda: clarify elementwise fusion aliasing check --- ggml/src/ggml-cuda/ggml-cuda.cu | 34 ++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2ad48b05e51a..fa57d8098dd8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3648,26 +3648,26 @@ static bool ggml_cuda_should_fuse_add_mul( return true; } -static bool ggml_cuda_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b) { - if (a == nullptr || b == nullptr || a->buffer == nullptr || b->buffer == nullptr) { - return false; - } - - const uintptr_t a_start = reinterpret_cast(a->data); - const uintptr_t b_start = reinterpret_cast(b->data); - const uintptr_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); - const uintptr_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); - - return a_start < b_end && b_start < a_end; -} - -static bool ggml_cuda_add_mul_memory_ok( +static bool ggml_cuda_check_elementwise_aliasing( const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * scale, const ggml_tensor * dst) { - const auto can_read_write_inplace = [dst](const ggml_tensor * src) { - return !ggml_cuda_tensors_overlap(dst, src) || ggml_are_same_layout(dst, src); + const auto tensors_overlap = [](const ggml_tensor * a, const ggml_tensor * b) { + if (a == nullptr || b == nullptr || a->buffer == nullptr || b->buffer == nullptr) { + return false; + } + + const uintptr_t a_start = reinterpret_cast(a->data); + const uintptr_t b_start = reinterpret_cast(b->data); + const uintptr_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + const uintptr_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + + return a_start < b_end && b_start < a_end; + }; + + const auto can_read_write_inplace = [dst, &tensors_overlap](const ggml_tensor * src) { + return !tensors_overlap(dst, src) || ggml_are_same_layout(dst, src); }; return can_read_write_inplace(src0) && @@ -4078,7 +4078,7 @@ static int ggml_cuda_try_fuse( const ggml_tensor * scale = nullptr; if (ggml_cuda_should_fuse_add_mul(cgraph->nodes[i], cgraph->nodes[i + 1], &src0, &src1, &scale)) { - if (ggml_cuda_add_mul_memory_ok(src0, src1, scale, cgraph->nodes[i + 1])) { + if (ggml_cuda_check_elementwise_aliasing(src0, src1, scale, cgraph->nodes[i + 1])) { ggml_cuda_op_add_mul_fused(*cuda_ctx, src0, src1, scale, cgraph->nodes[i + 1]); return 1; } From 50a2719d24216afa421a3294ab98d8ad6be2d5ad Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Fri, 3 Jul 2026 09:03:12 +0800 Subject: [PATCH 15/16] sycl: add rwkv lerp and rk kernels --- ggml/src/ggml-sycl/backend.hpp | 1 + ggml/src/ggml-sycl/ggml-sycl.cpp | 54 ++++++++++ ggml/src/ggml-sycl/rwkv.cpp | 167 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/rwkv.hpp | 10 ++ 4 files changed, 232 insertions(+) create mode 100644 ggml/src/ggml-sycl/rwkv.cpp create mode 100644 ggml/src/ggml-sycl/rwkv.hpp diff --git a/ggml/src/ggml-sycl/backend.hpp b/ggml/src/ggml-sycl/backend.hpp index 1f5a91272663..cfbf721017ae 100644 --- a/ggml/src/ggml-sycl/backend.hpp +++ b/ggml/src/ggml-sycl/backend.hpp @@ -38,6 +38,7 @@ #include "quants.hpp" #include "roll.hpp" #include "rope.hpp" +#include "rwkv.hpp" #include "set_rows.hpp" #include "ssm_conv.hpp" #include "softmax.hpp" diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 41449db665ec..70c08a6db446 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4948,6 +4948,12 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_TIMESTEP_EMBEDDING: ggml_sycl_op_timestep_embedding(ctx, dst); break; + case GGML_OP_RWKV_LERP: + ggml_sycl_op_rwkv_lerp(ctx, dst); + break; + case GGML_OP_RWKV_RK: + ggml_sycl_op_rwkv_rk(ctx, dst); + break; case GGML_OP_RWKV_WKV6: ggml_sycl_op_rwkv_wkv6(ctx, dst); break; @@ -5700,6 +5706,54 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons return true; case GGML_OP_LEAKY_RELU: case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_RWKV_LERP: + return op->type == GGML_TYPE_F32 && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + ggml_are_same_shape(op->src[0], op->src[1]) && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op->src[2]) && + ggml_is_contiguous(op) && + op->ne[0] == op->src[0]->ne[0] && + op->ne[1] == op->src[0]->ne[1] && + op->ne[2] == op->src[0]->ne[2] && + ggml_nelements(op) == ggml_nelements(op->src[0]) * op->ne[3] && + op->src[2]->ne[0] == op->ne[0] && + op->src[2]->ne[1] == 1 && + op->src[2]->ne[2] == 1 && + op->src[2]->ne[3] == op->ne[3] && + ggml_nelements(op->src[2]) == op->ne[0] * op->ne[3]; + case GGML_OP_RWKV_RK: { + const int64_t head_size = op->src[1]->ne[0]; + const int64_t H = op->src[1]->ne[1]; + const int64_t T = op->src[1]->ne[2]; + const int64_t C = head_size * H; + + return op->type == GGML_TYPE_F32 && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + op->src[3]->type == GGML_TYPE_F32 && + op->src[4]->type == GGML_TYPE_F32 && + (head_size == 64 || head_size == 128) && + ggml_are_same_shape(op->src[1], op->src[2]) && + ggml_are_same_shape(op->src[1], op->src[3]) && + ggml_are_same_shape(op->src[0], op) && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op->src[2]) && + ggml_is_contiguous(op->src[3]) && + ggml_is_contiguous(op->src[4]) && + ggml_is_contiguous(op) && + op->src[4]->ne[0] == head_size && + op->src[4]->ne[1] == H && + ggml_nelements(op->src[4]) == head_size * H && + op->ne[0] == C && + op->ne[1] == T && + ggml_nelements(op) == C * T; + } case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: case GGML_OP_GATED_LINEAR_ATTN: diff --git a/ggml/src/ggml-sycl/rwkv.cpp b/ggml/src/ggml-sycl/rwkv.cpp new file mode 100644 index 000000000000..bf221c7c4536 --- /dev/null +++ b/ggml/src/ggml-sycl/rwkv.cpp @@ -0,0 +1,167 @@ +#include "rwkv.hpp" + +#include + +static void rwkv_lerp_f32_kernel( + const float * x_prev, + const float * cur, + const float * weight, + float * dst, + const int64_t ne0, + const int64_t base_total, + const int64_t total, + sycl::id<1> id) { + const int64_t i = id[0]; + if (i >= total) { + return; + } + + const int64_t ibase = i % base_total; + const int64_t imix = i / base_total; + const int64_t iembd = ibase % ne0; + + const float c = cur[ibase]; + dst[i] = c + (x_prev[ibase] - c) * weight[imix * ne0 + iembd]; +} + +template +static void rwkv_rk_f32_kernel( + const float * cur, + const float * k, + const float * r, + const float * v, + const float * r_k, + float * dst, + const int64_t C, + const int64_t H, + sycl::nd_item<3> item, + float * shared_mem) { + const int64_t tid = item.get_local_id(2); + const int64_t row = item.get_group(2); + const int64_t h = row % H; + const int64_t t = row / H; + const int64_t off = t * C + h * head_size; + + shared_mem[tid] = k[off + tid] * r[off + tid] * r_k[h * head_size + tid]; + item.barrier(sycl::access::fence_space::local_space); + + for (int stride = head_size / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + shared_mem[tid] += shared_mem[tid + stride]; + } + item.barrier(sycl::access::fence_space::local_space); + } + + dst[off + tid] = cur[off + tid] + v[off + tid] * shared_mem[0]; +} + +void ggml_sycl_op_rwkv_lerp(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3); + + const ggml_tensor * x_prev = dst->src[0]; + const ggml_tensor * cur = dst->src[1]; + const ggml_tensor * weight = dst->src[2]; + + GGML_ASSERT(x_prev->type == GGML_TYPE_F32); + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(weight->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(x_prev, cur)); + GGML_ASSERT(ggml_is_contiguous(x_prev)); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(weight)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t base_total = ggml_nelements(x_prev); + const int64_t total = ggml_nelements(dst); + const int64_t n_mix = dst->ne[3]; + const int64_t ne0 = dst->ne[0]; + + GGML_ASSERT(dst->ne[0] == x_prev->ne[0]); + GGML_ASSERT(dst->ne[1] == x_prev->ne[1]); + GGML_ASSERT(dst->ne[2] == x_prev->ne[2]); + GGML_ASSERT(total == base_total * n_mix); + GGML_ASSERT(weight->ne[0] == dst->ne[0]); + GGML_ASSERT(weight->ne[1] == 1); + GGML_ASSERT(weight->ne[2] == 1); + GGML_ASSERT(weight->ne[3] == n_mix); + GGML_ASSERT(ggml_nelements(weight) == dst->ne[0] * n_mix); + + const float * x_prev_d = (const float *) x_prev->data; + const float * cur_d = (const float *) cur->data; + const float * weight_d = (const float *) weight->data; + float * dst_d = (float *) dst->data; + + dpct::queue_ptr stream = ctx.stream(); + stream->parallel_for(sycl::range<1>(total), [=](sycl::id<1> id) { + rwkv_lerp_f32_kernel(x_prev_d, cur_d, weight_d, dst_d, ne0, base_total, total, id); + }); +} + +void ggml_sycl_op_rwkv_rk(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/5); + + const ggml_tensor * cur = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * r = dst->src[2]; + const ggml_tensor * v = dst->src[3]; + const ggml_tensor * r_k = dst->src[4]; + + GGML_ASSERT(cur->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(r->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(r_k->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(cur)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r_k)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(k, r)); + GGML_ASSERT(ggml_are_same_shape(k, v)); + GGML_ASSERT(ggml_are_same_shape(cur, dst)); + + const int64_t head_size = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t T = k->ne[2]; + const int64_t C = head_size * H; + + GGML_ASSERT(head_size == 64 || head_size == 128); + GGML_ASSERT(r_k->ne[0] == head_size); + GGML_ASSERT(r_k->ne[1] == H); + GGML_ASSERT(ggml_nelements(r_k) == head_size * H); + GGML_ASSERT(dst->ne[0] == C); + GGML_ASSERT(dst->ne[1] == T); + GGML_ASSERT(ggml_nelements(dst) == C * T); + + const float * cur_d = (const float *) cur->data; + const float * k_d = (const float *) k->data; + const float * r_d = (const float *) r->data; + const float * v_d = (const float *) v->data; + const float * r_k_d = (const float *) r_k->data; + float * dst_d = (float *) dst->data; + + dpct::queue_ptr stream = ctx.stream(); + sycl::range<3> block_dims(1, 1, head_size); + sycl::range<3> grid_dims(1, 1, H * T); + + if (head_size == 64) { + stream->submit([&](sycl::handler & cgh) { + sycl::local_accessor shared_mem_acc(head_size, cgh); + cgh.parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item) { + rwkv_rk_f32_kernel<64>(cur_d, k_d, r_d, v_d, r_k_d, dst_d, C, H, item, + (float *) shared_mem_acc.get_multi_ptr().get()); + }); + }); + } else { + stream->submit([&](sycl::handler & cgh) { + sycl::local_accessor shared_mem_acc(head_size, cgh); + cgh.parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims), [=](sycl::nd_item<3> item) { + rwkv_rk_f32_kernel<128>(cur_d, k_d, r_d, v_d, r_k_d, dst_d, C, H, item, + (float *) shared_mem_acc.get_multi_ptr().get()); + }); + }); + } +} diff --git a/ggml/src/ggml-sycl/rwkv.hpp b/ggml/src/ggml-sycl/rwkv.hpp new file mode 100644 index 000000000000..ace38f766902 --- /dev/null +++ b/ggml/src/ggml-sycl/rwkv.hpp @@ -0,0 +1,10 @@ +#ifndef GGML_SYCL_RWKV_HPP +#define GGML_SYCL_RWKV_HPP + +#include "common.hpp" + +void ggml_sycl_op_rwkv_lerp(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +void ggml_sycl_op_rwkv_rk(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#endif // GGML_SYCL_RWKV_HPP From bf0fda2dfe02dc7eaa73cb953dce264dd496cfb3 Mon Sep 17 00:00:00 2001 From: Molly Sophia Date: Fri, 3 Jul 2026 09:06:20 +0800 Subject: [PATCH 16/16] cuda: rename rwkv fused ops file --- ggml/src/ggml-cuda/ggml-cuda.cu | 2 +- ggml/src/ggml-cuda/{fused-ops.cu => rwkv.cu} | 2 +- ggml/src/ggml-cuda/{fused-ops.cuh => rwkv.cuh} | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename ggml/src/ggml-cuda/{fused-ops.cu => rwkv.cu} (99%) rename ggml/src/ggml-cuda/{fused-ops.cuh => rwkv.cuh} (100%) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index fa57d8098dd8..1e6fe695aafc 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -25,7 +25,7 @@ #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" -#include "ggml-cuda/fused-ops.cuh" +#include "ggml-cuda/rwkv.cuh" #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" diff --git a/ggml/src/ggml-cuda/fused-ops.cu b/ggml/src/ggml-cuda/rwkv.cu similarity index 99% rename from ggml/src/ggml-cuda/fused-ops.cu rename to ggml/src/ggml-cuda/rwkv.cu index 74ae905eda45..2bc79b1acf1a 100644 --- a/ggml/src/ggml-cuda/fused-ops.cu +++ b/ggml/src/ggml-cuda/rwkv.cu @@ -1,4 +1,4 @@ -#include "fused-ops.cuh" +#include "rwkv.cuh" #include #include diff --git a/ggml/src/ggml-cuda/fused-ops.cuh b/ggml/src/ggml-cuda/rwkv.cuh similarity index 100% rename from ggml/src/ggml-cuda/fused-ops.cuh rename to ggml/src/ggml-cuda/rwkv.cuh