From 6647db9c27760044950fd6f99060456ae3d15df3 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 29 Jul 2026 07:36:53 +0300 Subject: [PATCH] DS4: streamline RoPE (#2198) * Use RoPE in-place to skip concatenating tensors * Remove some code duplication * Remove commented out code --- ggml/src/ggml-cuda/rope.cu | 248 ++++++++++++++++++++++++++++----- ggml/src/ggml.c | 43 ++++-- src/graphs/build_deepseek4.cpp | 117 +++------------- 3 files changed, 260 insertions(+), 148 deletions(-) diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu index a75d94080..439843428 100644 --- a/ggml/src/ggml-cuda/rope.cu +++ b/ggml/src/ggml-cuda/rope.cu @@ -41,7 +41,7 @@ template static __global__ void rope_norm( const T * x, T * dst, const int ne0, const int ne1, const int s1, const int s2, const int n_dims, const int32_t * pos, const float freq_scale, const float ext_factor, const float attn_factor, - const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors) { + const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors, bool is_flipped) { const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= ne0) { @@ -56,10 +56,18 @@ static __global__ void rope_norm( const int idst = row_dst*ne0 + i0; const int ix = channel_x*s2 + row_x*s1 + i0; - if (i0 >= n_dims) { - dst[idst + 0] = x[ix + 0]; - dst[idst + 1] = x[ix + 1]; + const int rope_offset = is_flipped ? ne0 - n_dims : 0; + if (i0 >= n_dims) { + if (dst != x) { + if (is_flipped) { + dst[idst + 0 - n_dims] = x[ix + 0 - n_dims]; + dst[idst + 1 - n_dims] = x[ix + 1 - n_dims]; + } else { + dst[idst + 0] = x[ix + 0]; + dst[idst + 1] = x[ix + 1]; + } + } return; } @@ -72,11 +80,45 @@ static __global__ void rope_norm( rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + const float x0 = x[ix + 0 + rope_offset]; + const float x1 = x[ix + 1 + rope_offset]; + + dst[idst + 0 + rope_offset] = x0*cos_theta - x1*sin_theta; + dst[idst + 1 + rope_offset] = x0*sin_theta + x1*cos_theta; +} + +template +static __global__ void rope_norm_inplace( + T * x, const int ne1, const int s1, const int s2, const int n_dims, + const int32_t * pos, const float freq_scale, const float ext_factor, const float attn_factor, + const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors) { + const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); + + if (i0 >= n_dims) { + return; + } + + const int row_dst = blockDim.x*blockIdx.x + threadIdx.x; + + const int row_x = row_dst % ne1; + const int channel_x = row_dst / ne1; + + const int ix = channel_x*s2 + row_x*s1 + i0; + + const float theta_base = pos[channel_x]*powf(theta_scale, i0/2.0f); + + const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + + float cos_theta; + float sin_theta; + + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + const float x0 = x[ix + 0]; const float x1 = x[ix + 1]; - dst[idst + 0] = x0*cos_theta - x1*sin_theta; - dst[idst + 1] = x0*sin_theta + x1*cos_theta; + x[ix + 0] = x0*cos_theta - x1*sin_theta; + x[ix + 1] = x0*sin_theta + x1*cos_theta; } static __global__ void rope_norm_fast(const float * src0, const float * src1, float * dst, int ne0, int ne1, int nelem, @@ -114,7 +156,7 @@ template static __global__ void rope_neox( const T * x, T * dst, const int ne0, const int ne1, const int s1, const int s2, const int n_dims, const int32_t * pos, const float freq_scale, const float ext_factor, const float attn_factor, - const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors) { + const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors, bool is_flipped) { const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= ne0) { @@ -126,16 +168,63 @@ static __global__ void rope_neox( const int row_x = row_dst % ne1; const int channel_x = row_dst / ne1; - const int idst = row_dst*ne0 + i0/2; - const int ix = channel_x*s2 + row_x*s1 + i0/2; + const int idst = row_dst*ne0; + const int ix = channel_x*s2 + row_x*s1; + + const int rope_offset = is_flipped ? ne0 - n_dims : 0; + + if (is_flipped) { + if (i0 < rope_offset) { + if (dst != x) { + dst[idst + i0 + 0] = x[ix + i0 + 0]; + dst[idst + i0 + 1] = x[ix + i0 + 1]; + } + return; + } + } else { + if (i0 >= n_dims) { + if (dst != x) { + dst[idst + i0 + 0] = x[ix + i0 + 0]; + dst[idst + i0 + 1] = x[ix + i0 + 1]; + } + return; + } + } + + const float theta_base = pos[channel_x]*powf(theta_scale, (i0 - rope_offset)/2.0f); + + const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + + float cos_theta; + float sin_theta; + + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0 - rope_offset, ext_factor, attn_factor, cos_theta, sin_theta); + + const float x0 = x[ix + i0/2 + rope_offset]; + const float x1 = x[ix + i0/2 + n_dims/2 + rope_offset]; + + dst[idst + i0/2 + rope_offset] = x0*cos_theta - x1*sin_theta; + dst[idst + i0/2 + n_dims/2 + rope_offset] = x0*sin_theta + x1*cos_theta; +} + +template +static __global__ void rope_neox_inplace( + T * x, const int ne1, const int s1, const int s2, const int n_dims, + const int32_t * pos, const float freq_scale, const float ext_factor, const float attn_factor, + const rope_corr_dims corr_dims, const float theta_scale, const float * freq_factors) { + const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= n_dims) { - dst[idst + i0/2 + 0] = x[ix + i0/2 + 0]; - dst[idst + i0/2 + 1] = x[ix + i0/2 + 1]; - return; } + const int row_dst = blockDim.x*blockIdx.x + threadIdx.x; + + const int row_x = row_dst % ne1; + const int channel_x = row_dst / ne1; + + const int ix = channel_x*s2 + row_x*s1; + const float theta_base = pos[channel_x]*powf(theta_scale, i0/2.0f); const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; @@ -145,11 +234,11 @@ static __global__ void rope_neox( rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims/2]; + const float x0 = x[ix + i0/2]; + const float x1 = x[ix + i0/2 + n_dims/2]; - dst[idst + 0] = x0*cos_theta - x1*sin_theta; - dst[idst + n_dims/2] = x0*sin_theta + x1*cos_theta; + x[ix + i0/2] = x0*cos_theta - x1*sin_theta; + x[ix + i0/2 + n_dims/2] = x0*sin_theta + x1*cos_theta; } static __global__ void rope_neox_fast(const float * src0, const float * src1, float * dst, int ne0, int ne1, int nelem, @@ -460,7 +549,7 @@ template static void rope_norm_cuda( const T * x, T * dst, const int ne0, const int ne1, const int s1, const int s2, const int n_dims, const int nr, const int32_t * pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, - const rope_corr_dims corr_dims, const float * freq_factors, cudaStream_t stream) { + const rope_corr_dims corr_dims, const float * freq_factors, bool is_flipped, cudaStream_t stream) { GGML_ASSERT(ne0 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); const int n_blocks_x = (ne0 + 2*CUDA_ROPE_BLOCK_SIZE - 1) / (2*CUDA_ROPE_BLOCK_SIZE); @@ -471,10 +560,33 @@ static void rope_norm_cuda( if (freq_factors == nullptr) { rope_norm<<>>( x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors); + attn_factor, corr_dims, theta_scale, freq_factors, is_flipped); } else { rope_norm<<>>( x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, is_flipped); + } +} + +template +static void rope_norm_cuda_inplace( + T * x, const int ne1, const int s1, const int s2, const int n_dims, const int nr, + const int32_t * pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, + const rope_corr_dims corr_dims, const float * freq_factors, cudaStream_t stream) { + GGML_ASSERT(n_dims % 2 == 0); + const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); + const int n_blocks_x = (n_dims + 2*CUDA_ROPE_BLOCK_SIZE - 1) / (2*CUDA_ROPE_BLOCK_SIZE); + const dim3 block_nums(nr, n_blocks_x, 1); + + const float theta_scale = powf(freq_base, -2.0f/n_dims); + + if (freq_factors == nullptr) { + rope_norm_inplace<<>>( + x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors); + } else { + rope_norm_inplace<<>>( + x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors); } } @@ -483,7 +595,7 @@ template static void rope_neox_cuda( const T * x, T * dst, const int ne0, const int ne1, const int s1, const int s2, const int n_dims, const int nr, const int32_t * pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, - const rope_corr_dims corr_dims, const float * freq_factors, cudaStream_t stream) { + const rope_corr_dims corr_dims, const float * freq_factors, bool is_flipped, cudaStream_t stream) { GGML_ASSERT(ne0 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); const int n_blocks_x = (ne0 + 2*CUDA_ROPE_BLOCK_SIZE - 1) / (2*CUDA_ROPE_BLOCK_SIZE); @@ -494,10 +606,33 @@ static void rope_neox_cuda( if (freq_factors == nullptr) { rope_neox<<>>( x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors); + attn_factor, corr_dims, theta_scale, freq_factors, is_flipped); } else { rope_neox<<>>( x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, is_flipped); + } +} + +template +static void rope_neox_cuda_inplace( + T * x, const int ne1, const int s1, const int s2, const int n_dims, const int nr, + const int32_t * pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, + const rope_corr_dims corr_dims, const float * freq_factors, cudaStream_t stream) { + GGML_ASSERT(n_dims % 2 == 0); + const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); + const int n_blocks_x = (n_dims + 2*CUDA_ROPE_BLOCK_SIZE - 1) / (2*CUDA_ROPE_BLOCK_SIZE); + const dim3 block_nums(nr, n_blocks_x, 1); + + const float theta_scale = powf(freq_base, -2.0f/n_dims); + + if (freq_factors == nullptr) { + rope_neox_inplace<<>>( + x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors); + } else { + rope_neox_inplace<<>>( + x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors); } } @@ -698,18 +833,35 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) rope_corr_dims corr_dims; ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v); + bool is_flipped = dst->op_params[15] == 1 && !is_vision; + //if (is_flipped && src0_d == dst_d) printf("Flipped in-place for %s\n", dst->name); + // compute if (is_neox) { - if (src0->type == GGML_TYPE_F32) { - rope_neox_cuda( - (const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, - freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); - } else if (src0->type == GGML_TYPE_F16) { - rope_neox_cuda( - (const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, - freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + if (src0->data == dst->data) { + if (src0->type == GGML_TYPE_F32) { + auto x = (float *)dst_d + (is_flipped ? ne00 - n_dims : 0); + rope_neox_cuda_inplace(x, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + } else if (src0->type == GGML_TYPE_F16) { + auto x = (half *)dst_d + (is_flipped ? ne00 - n_dims : 0); + rope_neox_cuda_inplace(x, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + } else { + GGML_ABORT("fatal error"); + } } else { - GGML_ABORT("fatal error"); + if (src0->type == GGML_TYPE_F32) { + rope_neox_cuda( + (const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, is_flipped, stream); + } else if (src0->type == GGML_TYPE_F16) { + rope_neox_cuda( + (const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, is_flipped, stream); + } else { + GGML_ABORT("fatal error"); + } } } else if (is_mrope && !is_vision) { if (src0->type == GGML_TYPE_F32) { @@ -736,16 +888,30 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) GGML_ABORT("fatal error"); } } else { - if (src0->type == GGML_TYPE_F32) { - rope_norm_cuda( - (const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, - freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); - } else if (src0->type == GGML_TYPE_F16) { - rope_norm_cuda( - (const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, - freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + if (src0->data == dst->data) { + if (src0->type == GGML_TYPE_F32) { + auto x = (float *)dst->data + (is_flipped ? ne00 - n_dims : 0); + rope_norm_cuda_inplace(x, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + } else if (src0->type == GGML_TYPE_F16) { + auto x = (half *)dst->data + (is_flipped ? ne00 - n_dims : 0); + rope_norm_cuda_inplace(x, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, stream); + } else { + GGML_ABORT("fatal error"); + } } else { - GGML_ABORT("fatal error"); + if (src0->type == GGML_TYPE_F32) { + rope_norm_cuda( + (const float *) src0_d, (float *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, is_flipped, stream); + } else if (src0->type == GGML_TYPE_F16) { + rope_norm_cuda( + (const half *) src0_d, (half *) dst_d, ne00, ne01, s01, s02, n_dims, nr, pos, freq_scale, + freq_base, ext_factor, attn_factor, corr_dims, freq_factors, is_flipped, stream); + } else { + GGML_ABORT("fatal error"); + } } } } @@ -822,6 +988,8 @@ void ggml_cuda_op_rope_cache_impl(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_ASSERT(n_dims == dst->ne[0]); } + GGML_ASSERT(dst->op_params[15] == 0); + const float * freq_factors = NULL; if (dst->src[1] != NULL) { GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F32); @@ -860,6 +1028,7 @@ void ggml_cuda_op_rope_fast(ggml_backend_cuda_context & ctx, ggml_tensor * dst) GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(src1->type == dst->type); + GGML_ASSERT(dst->op_params[15] == 0); const int64_t ne00 = src0->ne[0]; // head dims const int64_t ne01 = src0->ne[1]; // num heads @@ -900,6 +1069,7 @@ void ggml_cuda_op_rope_fast(ggml_backend_cuda_context & ctx, ggml_tensor * dst) bool ggml_cuda_op_fused_rope_fast(ggml_backend_cuda_context & ctx, ggml_tensor * dst1, ggml_tensor * dst2) { if (dst1->src[1] != dst2->src[1]) return false; + if (dst1->op_params[15] != 0 || dst2->op_params[15] != 0) return false; const ggml_tensor * src0_1 = dst1->src[0]; const ggml_tensor * src0_2 = dst2->src[0]; @@ -956,6 +1126,7 @@ bool ggml_cuda_op_fused_rope_fast(ggml_backend_cuda_context & ctx, ggml_tensor * bool ggml_cuda_op_fused_rms_rope_fast(ggml_backend_cuda_context & ctx, ggml_tensor * dst1, ggml_tensor * dst2) { if (dst1->src[1] != dst2->src[1]) return false; + if (dst1->op_params[15] != 0) return false; const auto rms_1 = dst1->src[0]; const auto rms_2 = dst2->src[0]; @@ -1372,6 +1543,7 @@ bool ggml_cuda_op_rope_rope_impl(ggml_backend_cuda_context & ctx, ggml_tensor * if (dst1->type != dst2->type) return false; if (dst1->src[0]->type != GGML_TYPE_F32 && dst1->src[0]->type != GGML_TYPE_F16) return false; if (dst1->src[0]->type != dst1->type) return false; + if (dst1->op_params[15] != 0 || dst2->op_params[15] != 0) return false; const int64_t ne00 = dst1->src[0]->ne[0]; const int64_t ne01_1 = dst1->src[0]->ne[1]; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c21d19a7e..94dd995e1 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -20806,6 +20806,7 @@ static void ggml_compute_forward_rope_f32( const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; // ggml_rope_multi, multimodal rotary position embedding const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope const bool is_vision = mode == GGML_ROPE_TYPE_VISION; + const bool is_inplace = src0->data == dst->data; if (is_mrope) { GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0); @@ -20829,6 +20830,9 @@ static void ggml_compute_forward_rope_f32( const int32_t * pos = (const int32_t *) src1->data; + const bool is_flipped = dst->op_params[15] == 1 && !is_vision && !is_mrope; + const int rope_offset = is_flipped ? ne0 - n_dims : 0; + for (int64_t i3 = 0; i3 < ne3; i3++) { // batch for (int64_t i2 = 0; i2 < ne2; i2++) { // seq-len @@ -20870,7 +20874,7 @@ static void ggml_compute_forward_rope_f32( } } else { for (int64_t i0 = 0; i0 < n_dims; i0 += 2) { - const int64_t ic = i0/2; + const int64_t ic = i0/2 + rope_offset; const float cos_theta = cache[i0 + 0]; const float sin_theta = cache[i0 + 1]; @@ -20889,9 +20893,10 @@ static void ggml_compute_forward_rope_f32( for (int64_t i0 = 0; i0 < n_dims; i0 += 2) { const float cos_theta = cache[i0 + 0]; const float sin_theta = cache[i0 + 1]; + const int ic = i0 + rope_offset; - const float * const src = (float *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); - float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + const float * const src = (float *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); + float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); const float x0 = src[0]; const float x1 = src[1]; @@ -20901,6 +20906,10 @@ static void ggml_compute_forward_rope_f32( } } + if (is_inplace) { + continue; + } + if (is_vision) { for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { const int64_t ic = i0/2; @@ -20919,12 +20928,22 @@ static void ggml_compute_forward_rope_f32( } } else { // fill the remain channels with data from src tensor - for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { - const float * const src = (float *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); - float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + if (is_flipped) { + for (int64_t i0 = 0; i0 < rope_offset; i0 += 2) { + const float * const src = (float *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); + float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); - dst_data[0] = src[0]; - dst_data[1] = src[1]; + dst_data[0] = src[0]; + dst_data[1] = src[1]; + } + } else { + for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { + const float * const src = (float *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); + float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + + dst_data[0] = src[0]; + dst_data[1] = src[1]; + } } } } @@ -20958,12 +20977,14 @@ static void ggml_compute_forward_rope_f16( memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int)*4); + const bool is_flipped = dst->op_params[15] != 0; + if (is_flipped) { + // TODO: implement it + GGML_ABORT("Flipped RoPE is not implemented for f16"); + } GGML_TENSOR_UNARY_OP_LOCALS - //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3); - //printf("n_past = %d, ne2 = %d\n", n_past, ne2); - GGML_ASSERT(nb0 == sizeof(ggml_fp16_t)); const int ith = params->ith; diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 3525736d0..d64357328 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -674,48 +674,7 @@ static ggml_tensor * build_hc_head( return llm.build_mhc_weighted_sum(x, pre, n_embd, hc); } -static ggml_tensor * build_hca_compressed_kv_from_state( - ggml_context * ctx0, - llm_build_context & llm, - ggml_tensor * kv_state, - ggml_tensor * score_state, - ggml_tensor * state_read_idxs, - ggml_tensor * comp_pos, - ggml_tensor * norm, - int64_t n_embd_head, - int il) { - const int64_t n_embd_head_rope = llm.hparams.n_rot; - const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; - const int64_t n_blocks = comp_pos ? comp_pos->ne[0] : 0; - - GGML_ASSERT(n_blocks > 0); - GGML_ASSERT(state_read_idxs != nullptr); - - auto comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, llama_context::dsv4_runtime::HCA_RATIO, 1); - - llm.cb(comp, "hca_comp_merge", il); - - comp = llm.llm_build_norm(ctx0, comp, llm.hparams, norm, nullptr, LLM_NORM_RMS, llm.cb, il); - llm.cb(comp, "hca_comp_norm", il); - - ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - 0); - ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head_nope)); - comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, llm.rope_type, llm.n_ctx_orig, - llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor, - dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow); - comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); - llm.cb(comp, "hca_comp_out", il); - - return comp; -} - -static ggml_tensor * build_overlap_compressed_kv_from_state( +static ggml_tensor * build_compressed_kv_from_state( ggml_context * ctx0, llm_build_context & llm, ggml_tensor * kv_state, @@ -728,31 +687,24 @@ static ggml_tensor * build_overlap_compressed_kv_from_state( int il, const char * tag) { const int64_t n_embd_head_rope = llm.hparams.n_rot; - const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; const int64_t n_blocks = comp_pos ? comp_pos->ne[0] : 0; GGML_ASSERT(n_blocks > 0); GGML_ASSERT(state_read_idxs != nullptr); - ggml_tensor * comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, ratio, 0); + int type = ratio == llama_context::dsv4_runtime::HCA_RATIO ? 1 : 0; + ggml_tensor * comp = ggml_ds4_comp(ctx0, kv_state, score_state, state_read_idxs, ratio, type); llm.cb(comp, tag, il); comp = llm.llm_build_norm(ctx0, comp, llm.hparams, norm, nullptr, LLM_NORM_RMS, llm.cb, il); llm.cb(comp, tag, il); - ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - 0); - ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head_nope)); - comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, llm.rope_type, llm.n_ctx_orig, + comp = ggml_reshape_3d(ctx0, comp, n_embd_head, 1, n_blocks); + comp = ggml_rope_ext_inplace(ctx0, comp, comp_pos, nullptr, n_embd_head_rope, llm.rope_type, llm.n_ctx_orig, llm.hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor, dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow); - comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); + comp->op_params[15] = 1; llm.cb(comp, tag, il); return comp; @@ -844,7 +796,6 @@ static ggml_tensor * dsv4_build_lid_top_k( const auto & layer = llm.model.layers[il]; const int64_t n_embd_indexer_head = hparams.indexer_head_size; const int64_t n_embd_indexer_head_rope = hparams.n_rot; - const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; const int64_t n_indexer_head = hparams.indexer_n_head; const int64_t n_tokens = cur->ne[1]; const int64_t n_lid = llm.lctx.dsv4.lid_plan.n_kv; @@ -859,19 +810,11 @@ static ggml_tensor * dsv4_build_lid_top_k( llm.cb(indexer_q, "lid_q", il); indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens); - ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - 0); - ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_embd_indexer_head_rope, + indexer_q = ggml_rope_ext_inplace(ctx0, indexer_q, inp_pos, nullptr, n_embd_indexer_head_rope, llm.rope_type, llm.n_ctx_orig, hparams.dsv4_compress_rope_base, llm.freq_scale, llm.ext_factor, dsv4_rope_attn_factor(llm.freq_scale, llm.ext_factor), llm.beta_fast, llm.beta_slow); - indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); + indexer_q->op_params[15] = 1; llm.cb(indexer_q, "indexer_q", il); GGML_ASSERT(indexer_q->ne[0] % hadamard_block == 0); indexer_q = ggml_hadamard(ctx0, indexer_q, hadamard_block); @@ -964,19 +907,13 @@ static void ds4_build_comp(ggml_tensor * cur, llm_build_context & llm, ggml_cont if (inputs.state_write_idxs != nullptr && plan.state_write_idxs.size() > 0) { ggml_tensor * source_kv = dsv4_concat_named(ctx0, cache_state, state_kv, 1, (tag + "_source_kv").c_str()); ggml_tensor * source_score = dsv4_concat_named(ctx0, cache_score, state_score, 1, (tag + "_source_score").c_str()); - ggml_tensor * comp = !is_hca ? build_overlap_compressed_kv_from_state(ctx0, llm, + auto ratio = is_hca ? llama_context::dsv4_runtime::HCA_RATIO : llama_context::dsv4_runtime::CSA_RATIO; + ggml_tensor * comp = build_compressed_kv_from_state(ctx0, llm, source_kv, source_score, inputs.state_read_idxs, inputs.state_write_pos, - norm, - llama_context::dsv4_runtime::CSA_RATIO, - head_size, il, - (tag + "_state_compress").c_str()) : - build_hca_compressed_kv_from_state(ctx0, llm, - source_kv, source_score, - inputs.state_read_idxs, - inputs.state_write_pos, - norm, head_size, il); + norm, ratio, head_size, il, + (tag + "_state_compress").c_str()); if (do_hadamard) { const int hadamard_block = llama_model::hadamard_size(head_size); GGML_ASSERT(hadamard_block > 0); @@ -1026,7 +963,6 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu const int64_t n_embd_head = hparams.n_embd_head_k(0); const int64_t n_embd_head_rope = hparams.n_rot; - const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; const int64_t hc = hparams.dsv4_hc_mult; const auto n_tokens = llm.n_tokens; @@ -1066,19 +1002,10 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu q = llm.llm_build_norm(ctx0, q, hparams, norm, nullptr, LLM_NORM_RMS, cb, il); cb(q, (tag + "_norm").c_str(), il); q = ggml_reshape_3d(ctx0, q, n_embd_head, nhead, n_tokens); - ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, nhead, n_tokens, - ggml_row_size(q->type, n_embd_head), - ggml_row_size(q->type, n_embd_head) * nhead, - 0); - ggml_tensor * q_rope = ggml_view_3d(ctx0, q, n_embd_head_rope, nhead, n_tokens, - ggml_row_size(q->type, n_embd_head), - ggml_row_size(q->type, n_embd_head) * nhead, - ggml_row_size(q->type, n_embd_head_nope)); - q_rope = ggml_rope_ext(ctx0, q_rope, inp_pos, nullptr, n_embd_head_rope, llm.rope_type, n_ctx_orig_l, + q = ggml_rope_ext_inplace(ctx0, q, inp_pos, nullptr, n_embd_head_rope, llm.rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); - cb(q_rope, (tag + "_rope").c_str(), il); - q = ggml_concat(ctx0, q_nope, q_rope, 0); - cb(q, tag.c_str(), il); + q->op_params[15] = 1; + cb(q, (tag + "_rope").c_str(), il); return q; }; @@ -1262,18 +1189,10 @@ static ggml_tensor * ds4_attention(ggml_cgraph * gf, ggml_context * ctx0, llm_bu ggml_build_forward_expand(gf, attn); attn = ggml_reshape_3d(ctx0, attn, n_embd_head, n_head, n_tokens); - ggml_tensor * attn_nope = ggml_view_3d(ctx0, attn, n_embd_head_nope, n_head, n_tokens, - ggml_row_size(attn->type, n_embd_head), - ggml_row_size(attn->type, n_embd_head) * n_head, - 0); - ggml_tensor * attn_pe = ggml_view_3d(ctx0, attn, n_embd_head_rope, n_head, n_tokens, - ggml_row_size(attn->type, n_embd_head), - ggml_row_size(attn->type, n_embd_head) * n_head, - ggml_row_size(attn->type, n_embd_head_nope)); - attn_pe = ggml_rope_back(ctx0, attn_pe, inp_pos, nullptr, n_embd_head_rope, llm.rope_type, n_ctx_orig_l, + attn = ggml_rope_ext_inplace(ctx0, attn, inp_pos, nullptr, n_embd_head_rope, llm.rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); - cb(attn_pe, "attn_derope", il); - attn = ggml_concat(ctx0, attn_nope, attn_pe, 0); + attn->op = GGML_OP_ROPE_BACK; + attn->op_params[15] = 1; cb(attn, "attn", il); const int64_t o_group_dim = layer.wo_a->ne[0];