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

This commit is contained in:
leejet authored and GitHub committed 2026-10-09 04:32:36 +08:00
1 parent e16d26a96a
commit 228c707fde
7 files changed
+36 -21

No files matched your search

+15 -3
View File
@@ -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];
@@ -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;