refactor: register CUDA RoPE kernels and fusion in sd.cpp

CMakeLists.txt

@@ -343,14 +343,19 @@ if (SD_CUDA)
     # Keep the driver stub on downstream link lines when no driver is installed.
     target_link_libraries(${SD_LIB} PUBLIC CUDA::cuda_driver)
     set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA)
-    if (NOT SD_USE_UPSTREAM_GGML AND GGML_CUDA_CUTLASS_PATH)
+    if (NOT SD_USE_UPSTREAM_GGML)
         enable_language(CUDA)
-        target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu)
-        set_source_files_properties(src/core/cuda/cute_attention.cu PROPERTIES
-            INCLUDE_DIRECTORIES "${GGML_CUDA_CUTLASS_PATH}/include"
-            COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr")
-        if (MSVC)
-            set_property(SOURCE src/core/cuda/cute_attention.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+        target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu)
+        set_source_files_properties(src/core/cuda/rope.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+        if (GGML_CUDA_CUTLASS_PATH)
+            target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu)
+            set_source_files_properties(src/core/cuda/cute_attention.cu PROPERTIES
+                INCLUDE_DIRECTORIES "${GGML_CUDA_CUTLASS_PATH}/include"
+                COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr")
+            if (MSVC)
+                set_property(SOURCE src/core/cuda/cute_attention.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+            endif()
+            set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
         endif()
         set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
         target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart)

docs/ggml_extensions.md

@@ -67,7 +67,11 @@ compute flags, external uses, output flags, and views; the initial implementatio
 accepts a conservative set of pure operators with only the final node exposed
 as the fused output. The rule remains responsible for mathematical equivalence,
 dtype, shape, layout, and hardware constraints. Internal views are accepted when
-their storage belongs to the fused subgraph; the final output must own storage.
+their storage belongs to the fused subgraph. The final output must own storage,
+except for a terminal `CPY`, which retains its original destination and write
+semantics. The extension must handle input/output storage overlap, including
+allocator reuse, with scratch space when necessary. Other in-place writes are
+not eligible for extension fusion.
 
 An optional `optimize` callback runs before matching and allocation. Use
 `ggml_ext_graph_delay` to move a producer range after independent nodes. The
@@ -111,8 +115,17 @@ recording service is not implemented yet.
 `sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
 `RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
 vector. Both interleaved and split-half RoPE, partial rotation, and strided input
-rows are supported. Model construction uses the custom operator on CPU and
-retains existing backend implementations elsewhere.
+rows are supported. Model construction uses the custom operator on CPU and CUDA
+and retains existing backend implementations elsewhere.
+
+`src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override,
+and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
+its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
+position-encoding strides are supported by the standalone kernel; RMSNorm fusion
+requires contiguous position encoding and weights. Copy fusion uses workspace
+when the destination overlaps an input. These registrations do not require
+CUTLASS. The old CUDA RoPE dispatch is removed; HIP and MUSA retain their existing
+implementations.
 
 `SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
 API is unchanged. INT8 ConvRot retains its existing packed activation layout and
@@ -130,6 +143,9 @@ variables are honored.
 The combined INT8 GEMM epilogue feeding FP16 Q directly into attention still needs
 to be externalized. The current extension can fuse RMSNorm/RoPE with attention,
 while other INT8 GEMM and standalone projection fusions retain their native
-implementations. Vulkan execution services and kernels also remain pending.
+implementations. Native INT8 projection matchers do not recognize custom RoPE
+yet, so those graphs currently run GEMM followed by the external RoPE/attention
+path. Restoring their combined epilogues is part of the remaining migration.
+Vulkan execution services and kernels also remain pending.
 This implementation does not yet externalize all model-specific kernels or add
 dynamic dtypes.

src/core/cuda/cute_attention.cu

@@ -44,13 +44,9 @@
 #include <mutex>
 #include <type_traits>
 #include <unordered_map>
+#include "core/cuda/rope.h"
 #include "ggml-ext.h"
 
-static __device__ __forceinline__ float sd_warp_sum(float v) {
-    for (int mask = 16; mask > 0; mask >>= 1)
-        v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask));
-    return v;
-}
 static __device__ __forceinline__ float sd_warp_max(float v) {
     for (int mask = 16; mask > 0; mask >>= 1)
         v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
@@ -349,75 +345,6 @@ namespace sd_fattn_cute {
     }
 }
 
-template <int D, bool interleaved, typename T>
-static __global__ void rms_norm_mul_rope_apply_fused(
-    const float* x,
-    const float* weight,
-    const float* pe,
-    T* dst,
-    int64_t heads,
-    int64_t tokens,
-    int64_t rows,
-    int64_t nb1,
-    int64_t nb2,
-    int64_t nb3,
-    int rot_half,
-    float eps,
-    float output_scale) {
-    const int lane    = threadIdx.x;
-    const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
-    if (row >= rows) {
-        return;
-    }
-    const int64_t head  = row % heads;
-    const int64_t token = (row / heads) % tokens;
-    const int64_t batch = row / (heads * tokens);
-    const float* x_row  = x + head * nb1 + token * nb2 + batch * nb3;
-    float v[D / 32], sums[8] = {};
-#pragma unroll
-    for (int j = 0; j < D / 32; ++j) {
-        v[j]    = x_row[j * 32 + lane];
-        sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
-    }
-    // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
-    const float sum = __fadd_rn(
-        __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
-        __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
-    const float scale = rsqrtf(sum / float(D) + eps);
-    T* out            = dst + ((batch * heads + head) * tokens + token) * D;
-#pragma unroll
-    for (int j = 0; j < D / 32; ++j) {
-        const int col     = j * 32 + lane;
-        const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
-        float rotated     = value;
-        if constexpr (interleaved) {
-            const float other = __shfl_xor_sync(0xffffffff, value, 1, 32);
-            if (col < 2 * rot_half) {
-                const float* m      = pe + 4 * (token * rot_half + col / 2);
-                const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
-                rotated             = lane % 2 == 0
-                                          ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
-                                          : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
-            }
-        } else {
-            if (col < 2 * rot_half) {
-                const int partner   = col < rot_half ? col + rot_half : col - rot_half;
-                const float other   = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
-                const float* m      = pe + 4 * (token * rot_half + col % rot_half);
-                const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
-                rotated             = col < rot_half
-                                          ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
-                                          : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
-            }
-        }
-        if constexpr (std::is_same<T, half>::value) {
-            out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
-        } else {
-            out[col] = rotated;
-        }
-    }
-}
-
 static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
     a = index == 3 ? -a : a;
     b = index == 2 ? -b : b;
@@ -610,21 +537,6 @@ static bool fusion_enabled() {
     return !disabled;
 }
 
-static bool query_fusion(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
-    if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || rope->op != GGML_OP_ROPE_APPLY ||
-        mul->src[0] != norm || rope->src[0] != mul)
-        return false;
-    const auto* x      = norm->src[0];
-    const auto* weight = mul->src[1];
-    const auto* pe     = rope->src[1];
-    return x && weight && pe && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
-           x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
-           norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && rope->type == GGML_TYPE_F32 &&
-           ggml_are_same_shape(norm, mul) && weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
-           pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] > 0 && pe->ne[2] <= x->ne[0] / 2 && pe->ne[3] == x->ne[2] &&
-           x->nb[0] == sizeof(float) && ggml_is_contiguous(weight) && ggml_is_contiguous(pe) && ggml_is_contiguous(rope);
-}
-
 static const ggml_tensor* graph_node(const ggml_cgraph* graph, int index) {
     return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
 }
@@ -668,7 +580,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
     int offset = 0;
     if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
         if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
-            !query_fusion(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
+            !sd_cuda_rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
             graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
             return 0;
         offset = 3;
@@ -688,7 +600,7 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
             continue;
         for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
             if (attn->src[0] == graph_node(graph, j + 2) &&
-                query_fusion(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
+                sd_cuda_rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
                 ggml_ext_graph_delay(graph, j, 3, i);
                 break;
             }
@@ -741,30 +653,6 @@ static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const
     return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
 }
 
-template <bool Interleaved>
-static void query_launch(cudaStream_t stream, ggml_tensor* const* nodes, half* output, float scale) {
-    const auto* x      = nodes[0]->src[0];
-    const auto* pe     = nodes[2]->src[1];
-    const int64_t rows = ggml_nrows(x);
-#define SD_QUERY(D)                                                                                                                     \
-    case D:                                                                                                                             \
-        rms_norm_mul_rope_apply_fused<D, Interleaved, half><<<(rows + 3) / 4, dim3(32, 4), 0, stream>>>(                                \
-            static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data), static_cast<const float*>(pe->data), \
-            output, x->ne[1], x->ne[2], rows, x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
-        break
-    switch (x->ne[0]) {
-        SD_QUERY(32);
-        SD_QUERY(64);
-        SD_QUERY(96);
-        SD_QUERY(128);
-        SD_QUERY(160);
-        SD_QUERY(192);
-        SD_QUERY(224);
-        SD_QUERY(256);
-    }
-#undef SD_QUERY
-}
-
 static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
     const auto& plan = *static_cast<AttentionPlan*>(state);
     const auto* attn = nodes[plan.index];
@@ -784,10 +672,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
         const auto* src   = attn->src[i];
         const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
         if (i == 0 && plan.index == 3) {
-            if (nodes[2]->op_params[0])
-                query_launch<true>(stream, nodes, work[0], scale);
-            else
-                query_launch<false>(stream, nodes, work[0], scale);
+            if (sd_cuda_rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
+                return GGML_STATUS_FAILED;
         } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
             inputs[i] = static_cast<const half*>(src->data);
         } else {

src/core/cuda/rope.cu

@@ -0,0 +1,305 @@
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <cuda_fp16.h>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <type_traits>
+
+static __device__ __forceinline__ float sd_warp_sum(float v) {
+    for (int mask = 16; mask > 0; mask >>= 1)
+        v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask));
+    return v;
+}
+template <int D, bool interleaved, typename T>
+static __global__ void rms_norm_mul_rope_apply_fused(
+    const float* x,
+    const float* weight,
+    const float* pe,
+    T* dst,
+    int64_t heads,
+    int64_t tokens,
+    int64_t rows,
+    int64_t nb1,
+    int64_t nb2,
+    int64_t nb3,
+    int rot_half,
+    float eps,
+    float output_scale) {
+    const int lane    = threadIdx.x;
+    const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+    if (row >= rows) {
+        return;
+    }
+    const int64_t head  = row % heads;
+    const int64_t token = (row / heads) % tokens;
+    const int64_t batch = row / (heads * tokens);
+    const float* x_row  = x + head * nb1 + token * nb2 + batch * nb3;
+    float v[D / 32], sums[8] = {};
+#pragma unroll
+    for (int j = 0; j < D / 32; ++j) {
+        v[j]    = x_row[j * 32 + lane];
+        sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
+    }
+    // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
+    const float sum = __fadd_rn(
+        __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
+        __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
+    const float scale = rsqrtf(sum / float(D) + eps);
+    T* out            = dst + ((batch * heads + head) * tokens + token) * D;
+#pragma unroll
+    for (int j = 0; j < D / 32; ++j) {
+        const int col     = j * 32 + lane;
+        const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
+        float rotated     = value;
+        if constexpr (interleaved) {
+            const float other = __shfl_xor_sync(0xffffffff, value, 1, 32);
+            if (col < 2 * rot_half) {
+                const float* m      = pe + 4 * (token * rot_half + col / 2);
+                const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                rotated             = lane % 2 == 0
+                                          ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
+                                          : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
+            }
+        } else {
+            if (col < 2 * rot_half) {
+                const int partner   = col < rot_half ? col + rot_half : col - rot_half;
+                const float other   = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
+                const float* m      = pe + 4 * (token * rot_half + col % rot_half);
+                const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                rotated             = col < rot_half
+                                          ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
+                                          : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
+            }
+        }
+        if constexpr (std::is_same<T, half>::value) {
+            out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
+        } else {
+            out[col] = rotated;
+        }
+    }
+}
+
+struct RopeApplyStrides {
+    int64_t x[4];
+    int64_t pe[4];
+};
+
+template <bool interleaved>
+static __global__ void rope_apply_f32(
+    const float* x,
+    const float* pe,
+    float* dst,
+    int64_t half_dim,
+    int64_t rot_half,
+    int64_t heads,
+    int64_t tokens,
+    int64_t pairs,
+    RopeApplyStrides strides) {
+    const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (index >= pairs) {
+        return;
+    }
+    const int64_t j     = index % half_dim;
+    const int64_t row   = index / half_dim;
+    const int64_t token = row % tokens;
+    const int64_t head  = (row / tokens) % heads;
+    const int64_t batch = row / (tokens * heads);
+    const int64_t i0    = interleaved || j >= rot_half ? 2 * j : j;
+    const int64_t i1    = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
+    const float* x_row  = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
+    const float a       = x_row[i0 * strides.x[0]];
+    const float b       = x_row[i1 * strides.x[0]];
+    float* out          = dst + row * half_dim * 2;
+    if (j >= rot_half) {
+        out[i0] = a;
+        out[i1] = b;
+        return;
+    }
+    const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
+    const float c       = matrix[0];
+    const float s       = matrix[strides.pe[1]];
+    // Match the separate mul/add rounding of the unfused graph.
+    if (interleaved) {
+        out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
+        out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
+    } else {
+        out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
+        out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
+    }
+}
+
+static bool rope_interleaved(const ggml_tensor* rope) {
+    int32_t value;
+    std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
+    return value != 0;
+}
+
+static float param_f32(const ggml_tensor* node, int index) {
+    float value;
+    std::memcpy(&value, node->op_params + index, sizeof(value));
+    return value;
+}
+
+bool sd_cuda_rope_is_node(const ggml_tensor* node) {
+    return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
+}
+
+static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
+    const auto* x  = rope->src[0];
+    const auto* pe = rope->src[1];
+    if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
+        x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+        pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
+        rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
+        (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
+        return false;
+    }
+    for (int i = 0; i < 4; ++i) {
+        if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
+            return false;
+    }
+    return true;
+}
+
+bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+    if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !sd_cuda_rope_is_node(rope) ||
+        mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
+        return false;
+    const auto* x      = norm->src[0];
+    const auto* weight = mul->src[1];
+    return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+           x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
+           ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
+           weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
+           x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
+           x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
+           (ggml_nrows(x) + 3) / 4 <= INT_MAX;
+}
+
+template <typename T, bool Interleaved>
+static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) {
+    const auto* x      = nodes[0]->src[0];
+    const auto* pe     = nodes[2]->src[1];
+    const int64_t rows = ggml_nrows(x);
+#define SD_ROPE(D)                                                                                                               \
+    case D:                                                                                                                      \
+        rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
+            static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),                               \
+            static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows,                                               \
+            x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale);                            \
+        break
+    switch (x->ne[0]) {
+        SD_ROPE(32);
+        SD_ROPE(64);
+        SD_ROPE(96);
+        SD_ROPE(128);
+        SD_ROPE(160);
+        SD_ROPE(192);
+        SD_ROPE(224);
+        SD_ROPE(256);
+        default:
+            return cudaErrorInvalidValue;
+    }
+#undef SD_ROPE
+    return cudaGetLastError();
+}
+
+cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
+    if (fp16) {
+        return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
+                                          : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
+    }
+    return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
+                                      : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
+}
+
+static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+    static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
+    if (disabled)
+        return 0;
+    auto* g = const_cast<ggml_cgraph*>(graph);
+    if (first + 2 >= ggml_graph_n_nodes(g) || !sd_cuda_rope_fusion_supports(
+                                                  ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
+        return 0;
+    if (first + 4 < ggml_graph_n_nodes(g)) {
+        const auto* rope  = ggml_graph_node(g, first + 2);
+        const auto* scale = ggml_graph_node(g, first + 3);
+        const auto* copy  = ggml_graph_node(g, first + 4);
+        if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
+            scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
+            ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+            return 5;
+    }
+    return 3;
+}
+
+static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+    *workspace = ggml_nbytes(nodes[count - 1]);
+    return GGML_STATUS_SUCCESS;
+}
+
+static void rope_destroy(void*, void*) {}
+
+static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+    const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+    return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+}
+
+static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+    auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+    auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    auto stream      = static_cast<cudaStream_t>(context.stream);
+    const auto* rope = nodes[count == 1 ? 0 : 2];
+    const auto* x    = count == 1 ? rope->src[0] : nodes[0]->src[0];
+    auto* dst        = nodes[count - 1];
+    const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
+    void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+    cudaError_t error;
+    if (count > 1) {
+        error = sd_cuda_rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+    } else {
+        RopeApplyStrides strides;
+        const auto* pe = rope->src[1];
+        for (int i = 0; i < 4; ++i) {
+            strides.x[i]  = x->nb[i] / 4;
+            strides.pe[i] = pe->nb[i] / 4;
+        }
+        const int64_t pairs = ggml_nelements(x) / 2;
+        auto kernel         = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
+        kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
+                                                                                   static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
+        error = cudaGetLastError();
+    }
+    if (error == cudaSuccess && alias)
+        error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+    return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) {
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.rope_apply.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_rope_op();
+    desc.priority             = 100;
+    desc.supports             = rope_supports;
+    desc.prepare              = rope_prepare;
+    desc.destroy              = rope_destroy;
+    desc.execute              = rope_execute;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    desc.name      = "sd.rope_apply.cuda.override.v1";
+    desc.op        = GGML_OP_ROPE_APPLY;
+    desc.custom_op = nullptr;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    desc.name     = "sd.rms_norm_mul_rope.cuda.v1";
+    desc.op       = GGML_OP_RMS_NORM;
+    desc.priority = 200;
+    desc.supports = nullptr;
+    desc.match    = rope_match;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+}

src/core/cuda/rope.h

@@ -0,0 +1,12 @@
+#ifndef __SD_CORE_CUDA_ROPE_H__
+#define __SD_CORE_CUDA_ROPE_H__
+
+#include <cuda_runtime.h>
+#include "ggml-ext.h"
+
+bool sd_cuda_rope_is_node(const ggml_tensor* node);
+bool sd_cuda_rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+cudaError_t sd_cuda_rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
+void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
+
+#endif

src/core/ggml_extensions.cpp

@@ -9,6 +9,9 @@
 #include <unordered_set>
 
 #ifdef SD_CUDA_EXTENSIONS
+void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
+#endif
+#ifdef SD_CUDA_CUTE_ATTENTION
 void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
 #endif
 
@@ -31,7 +34,7 @@ static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_
     return true;
 }
 
-static ggml_ext_op_t rope_op() {
+ggml_ext_op_t sd_ggml_rope_op() {
     static const ggml_ext_op_t op = [] {
         const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.rope_apply.v1", sizeof(int32_t), rope_infer};
         auto result                 = ggml_ext_register_op(&desc);
@@ -130,7 +133,7 @@ static int rope_fusion_match(ggml_backend_dev_t device, const ggml_cgraph* graph
     if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
         norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
         mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
-        (rope->op != GGML_OP_ROPE_APPLY && ggml_ext_op_get(rope) != rope_op()) ||
+        (rope->op != GGML_OP_ROPE_APPLY && ggml_ext_op_get(rope) != sd_ggml_rope_op()) ||
         !rope_supports(device, rope, nullptr) || !ggml_are_same_shape(norm, mul)) {
         return 0;
     }
@@ -174,6 +177,9 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
     }
     if (cuda) {
 #ifdef SD_CUDA_EXTENSIONS
+        sd_ggml_register_cuda_rope(device);
+#endif
+#ifdef SD_CUDA_CUTE_ATTENTION
         sd_ggml_register_cuda_attention(device);
 #endif
         return;
@@ -182,7 +188,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
     desc.abi_version          = GGML_EXT_ABI_VERSION;
     desc.name                 = "sd.rope_apply.cpu.v1";
     desc.op                   = GGML_OP_CUSTOM;
-    desc.custom_op            = rope_op();
+    desc.custom_op            = sd_ggml_rope_op();
     desc.priority             = 100;
     desc.supports             = rope_supports;
     desc.execute              = rope_execute;
@@ -205,10 +211,14 @@ ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_
     sd_ggml_register_extensions(backend);
     if (backend) {
         auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
-        if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
+        if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")
+#ifdef SD_CUDA_EXTENSIONS
+            || ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1")
+#endif
+        ) {
             ggml_tensor* inputs[] = {x, pe};
             int32_t attributes    = interleaved ? 1 : 0;
-            return ggml_ext_op_create(ctx, rope_op(), inputs, 2, &attributes);
+            return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
         }
     }
     return ggml_rope_apply(ctx, x, pe, interleaved);

src/core/ggml_extensions.h

@@ -6,6 +6,7 @@
 void sd_ggml_register_extensions(ggml_backend_t backend);
 
 #ifndef SD_USE_UPSTREAM_GGML
+const struct ggml_ext_op* sd_ggml_rope_op();
 ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
 #endif