mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 03:49:09 +00:00
perf: use ggml rope apply op on supported backends (#2113)
This commit is contained in:
1 parent
e16d26a96a
commit
228c707fde
7 files changed
+36
-21
No files matched your search
@@ -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;
|
||||
|
||||
Reference in new issue
Block a user