Skip to content

perf: use ggml rope apply op on supported backends #2113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leejet merged 1 commit into from
Oct 8, 2026
Merged

perf: use ggml rope apply op on supported backends #2113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ggml
Open in desktop
Submodule ggml updated 14 files
+12 −0 include/ggml.h
+5 −0 src/ggml-cpu/ggml-cpu.c
+3 −0 src/ggml-cpu/ggml-cpu.cpp
+42 −0 src/ggml-cpu/ops.cpp
+1 −0 src/ggml-cpu/ops.h
+6 −0 src/ggml-cuda/ggml-cuda.cu
+63 −0 src/ggml-cuda/rope.cu
+2 −0 src/ggml-cuda/rope.cuh
+15 −0 src/ggml-vulkan/ggml-vulkan-push-constants.h
+1 −0 src/ggml-vulkan/ggml-vulkan-types.h
+51 −0 src/ggml-vulkan/ggml-vulkan.cpp
+53 −0 src/ggml-vulkan/vulkan-shaders/rope_apply.comp
+2 −0 src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+24 −2 src/ggml.c
18 changes: 15 additions & 3 deletions src/model/common/rope.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1067,10 +1067,22 @@ namespace Rope {
return result;
}

__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* apply_rope(GGMLRunnerContext* runner_ctx,
ggml_tensor* x,
ggml_tensor* pe,
bool rope_interleaved = true) {
auto ctx = runner_ctx->ggml_ctx;
#ifndef SD_USE_UPSTREAM_GGML
auto backend = runner_ctx->backend;
if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 &&
pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) {
auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved);
if (ggml_backend_supports_op(backend, out)) {
return out;
}
}
#endif
// x: [N, L, n_head, d_head]
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
int64_t d_head = x->ne[0];
Expand Down Expand Up @@ -1126,8 +1138,8 @@ namespace Rope {
// return: [N, L, n_head*d_head]
int64_t n_head = q->ne[1];

q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head]
k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head]
q = apply_rope(ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head]
k = apply_rope(ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head]

auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head]
return x;
Expand Down
4 changes: 2 additions & 2 deletions src/model/diffusion/anima.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,8 +235,8 @@ namespace Anima {
if (pe_k == nullptr) {
pe_k = pe_q;
}
auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
auto q_rope = Rope::apply_rope(ctx, q4, pe_q, false);
auto k_rope = Rope::apply_rope(ctx, k4, pe_k, false);
attn_out = ggml_ext_attention_ext(ctx,
q_rope,
k_rope,
Expand Down
11 changes: 6 additions & 5 deletions src/model/diffusion/ltxv.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -548,21 +548,22 @@ namespace LTXV {
return build_rope_matrix_from_frequencies(freqs, dim);
}

__STATIC_INLINE__ ggml_tensor* apply_hidden_rope(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* apply_hidden_rope(GGMLRunnerContext* runner_ctx,
ggml_tensor* x,
ggml_tensor* pe,
int64_t heads,
int64_t dim_head,
bool rope_interleaved) {
auto ctx = runner_ctx->ggml_ctx;
GGML_ASSERT(x->ne[0] == heads * dim_head);
auto x4 = ggml_reshape_4d(ctx, x, dim_head, heads, x->ne[1], x->ne[2]);
if (pe != nullptr && pe->ne[3] == x->ne[1] * heads) {
auto x_flat = ggml_reshape_4d(ctx, x4, dim_head, 1, x->ne[1] * heads, x->ne[2]);
auto out_flat = Rope::apply_rope(ctx, x_flat, pe, rope_interleaved);
auto out_flat = Rope::apply_rope(runner_ctx, x_flat, pe, rope_interleaved);
auto out4 = ggml_reshape_4d(ctx, out_flat, dim_head, heads, x->ne[1], x->ne[2]);
return ggml_reshape_3d(ctx, out4, heads * dim_head, x->ne[1], x->ne[2]);
}
return Rope::apply_rope(ctx, x4, pe, rope_interleaved);
return Rope::apply_rope(runner_ctx, x4, pe, rope_interleaved);
}

struct TimestepEmbedder : public GGMLBlock {
Expand Down Expand Up @@ -705,8 +706,8 @@ namespace LTXV {
if (k_pe == nullptr) {
k_pe = pe;
}
q = apply_hidden_rope(ctx->ggml_ctx, q, pe, heads, dim_head, rope_interleaved);
k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
q = apply_hidden_rope(ctx, q, pe, heads, dim_head, rope_interleaved);
k = apply_hidden_rope(ctx, k, k_pe, heads, dim_head, rope_interleaved);
}

auto out = ggml_ext_attention_ext(ctx,
Expand Down
9 changes: 5 additions & 4 deletions src/model/diffusion/minimax_h3.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,12 +160,13 @@ namespace MiniMaxH3 {
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
}

static ggml_tensor* apply_partial_rope(ggml_context* ctx,
static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto ctx = runner_ctx->ggml_ctx;
int64_t rot_dim = pe->ne[2] * 2;
GGML_ASSERT(rot_dim <= x->ne[0]);
auto rotated = Rope::apply_rope(ctx,
auto rotated = Rope::apply_rope(runner_ctx,
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
pe,
false);
Expand Down Expand Up @@ -209,8 +210,8 @@ namespace MiniMaxH3 {
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
if (pe != nullptr) {
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
q = apply_partial_rope(ctx, q, pe);
k = apply_partial_rope(ctx, k, pe);
} else {
q = attention_layout(ctx->ggml_ctx, q);
k = attention_layout(ctx->ggml_ctx, k);
Expand Down
4 changes: 2 additions & 2 deletions src/model/diffusion/qwen_image_2_1.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,8 @@ namespace Qwen {
auto v = project("to_v");
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
q = Rope::apply_rope(ctx, q, pe);
k = Rope::apply_rope(ctx, k, pe);
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
// Preserve query-first attention evaluation while writing each layer's
// prefix before its full-sequence K/V can accumulate across layers.
Expand Down
9 changes: 5 additions & 4 deletions src/model/vae/minimax_h3_vae.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,11 +228,12 @@ namespace MiniMaxH3VAE {
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
}

static ggml_tensor* apply_partial_rope(ggml_context* ctx,
static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto ctx = runner_ctx->ggml_ctx;
int64_t rot_dim = pe->ne[2] * 2;
auto rotated = Rope::apply_rope(ctx,
auto rotated = Rope::apply_rope(runner_ctx,
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
pe,
false);
Expand Down Expand Up @@ -289,8 +290,8 @@ namespace MiniMaxH3VAE {
batch_size);
q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
q = apply_partial_rope(ctx, q, pe);
k = apply_partial_rope(ctx, k, pe);
auto out = ggml_ext_attention_ext(ctx,
q,
k,
Expand Down
Loading