DS4: streamline RoPE (#2198)

* Use RoPE in-place to skip concatenating tensors * Remove some code duplication * Remove commented out code
2026-08-12 22:29:39 +04:00 · 2026-07-29 07:36:53 +03:00
parent b054a8b983
commit 6647db9c27
3 changed files with 260 additions and 148 deletions
@@ -41,7 +41,7 @@ template<bool forward, bool has_ff, typename T>
 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<forward>(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<bool forward, bool has_ff, typename T>
+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<forward>(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<bool forward, bool has_ff, typename T>
 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<forward>(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<bool forward, bool has_ff, typename T>
+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<forward>(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<bool forward, typename T>
 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<forward, false><<<block_nums, block_dims, 0, stream>>>(
            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<forward, true><<<block_nums, block_dims, 0, stream>>>(
            x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, is_flipped);
+    }
+}
+
+template<bool forward, typename T>
+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<forward, false><<<block_nums, block_dims, 0, stream>>>(
+            x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors);
+    } else {
+        rope_norm_inplace<forward, true><<<block_nums, block_dims, 0, stream>>>(
+            x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor,
            attn_factor, corr_dims, theta_scale, freq_factors);
    }
 }
@@ -483,7 +595,7 @@ template<bool forward, typename T>
 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<forward, false, T><<<block_nums, block_dims, 0, stream>>>(
            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<forward, true, T><<<block_nums, block_dims, 0, stream>>>(
            x, dst, ne0, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors, is_flipped);
+    }
+}
+
+template<bool forward, typename T>
+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<forward, false, T><<<block_nums, block_dims, 0, stream>>>(
+            x, ne1, s1, s2, n_dims, pos, freq_scale, ext_factor,
+            attn_factor, corr_dims, theta_scale, freq_factors);
+    } else {
+        rope_neox_inplace<forward, true, T><<<block_nums, block_dims, 0, stream>>>(
+            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<forward>(
-                (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<forward>(
-                (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<forward>(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<forward>(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<forward>(
+                        (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<forward>(
+                        (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<forward>(
-                (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<forward>(
-                (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<forward>(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<forward>(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<forward>(
+                        (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<forward>(
+                        (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];
@@ -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(&sections,    (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;
@@ -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];