mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +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
+1
-1
Submodule ggml updated: 89c4413f5d...66c1a28760.
@@ -1067,10 +1067,22 @@ namespace Rope {
|
|||||||
return result;
|
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* x,
|
||||||
ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
bool rope_interleaved = true) {
|
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]
|
// x: [N, L, n_head, d_head]
|
||||||
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
|
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
|
||||||
int64_t d_head = x->ne[0];
|
int64_t d_head = x->ne[0];
|
||||||
@@ -1126,8 +1138,8 @@ namespace Rope {
|
|||||||
// return: [N, L, n_head*d_head]
|
// return: [N, L, n_head*d_head]
|
||||||
int64_t n_head = q->ne[1];
|
int64_t n_head = q->ne[1];
|
||||||
|
|
||||||
q = apply_rope(ctx->ggml_ctx, q, 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->ggml_ctx, k, 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]
|
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;
|
return x;
|
||||||
|
|||||||
@@ -235,8 +235,8 @@ namespace Anima {
|
|||||||
if (pe_k == nullptr) {
|
if (pe_k == nullptr) {
|
||||||
pe_k = pe_q;
|
pe_k = pe_q;
|
||||||
}
|
}
|
||||||
auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
|
auto q_rope = Rope::apply_rope(ctx, q4, pe_q, false);
|
||||||
auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
|
auto k_rope = Rope::apply_rope(ctx, k4, pe_k, false);
|
||||||
attn_out = ggml_ext_attention_ext(ctx,
|
attn_out = ggml_ext_attention_ext(ctx,
|
||||||
q_rope,
|
q_rope,
|
||||||
k_rope,
|
k_rope,
|
||||||
|
|||||||
@@ -548,21 +548,22 @@ namespace LTXV {
|
|||||||
return build_rope_matrix_from_frequencies(freqs, dim);
|
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* x,
|
||||||
ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
int64_t heads,
|
int64_t heads,
|
||||||
int64_t dim_head,
|
int64_t dim_head,
|
||||||
bool rope_interleaved) {
|
bool rope_interleaved) {
|
||||||
|
auto ctx = runner_ctx->ggml_ctx;
|
||||||
GGML_ASSERT(x->ne[0] == heads * dim_head);
|
GGML_ASSERT(x->ne[0] == heads * dim_head);
|
||||||
auto x4 = ggml_reshape_4d(ctx, x, dim_head, heads, x->ne[1], x->ne[2]);
|
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) {
|
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 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]);
|
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 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 {
|
struct TimestepEmbedder : public GGMLBlock {
|
||||||
@@ -705,8 +706,8 @@ namespace LTXV {
|
|||||||
if (k_pe == nullptr) {
|
if (k_pe == nullptr) {
|
||||||
k_pe = pe;
|
k_pe = pe;
|
||||||
}
|
}
|
||||||
q = apply_hidden_rope(ctx->ggml_ctx, q, pe, heads, dim_head, rope_interleaved);
|
q = apply_hidden_rope(ctx, q, pe, heads, dim_head, rope_interleaved);
|
||||||
k = apply_hidden_rope(ctx->ggml_ctx, k, k_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,
|
auto out = ggml_ext_attention_ext(ctx,
|
||||||
|
|||||||
@@ -160,12 +160,13 @@ namespace MiniMaxH3 {
|
|||||||
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
|
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* x,
|
||||||
ggml_tensor* pe) {
|
ggml_tensor* pe) {
|
||||||
|
auto ctx = runner_ctx->ggml_ctx;
|
||||||
int64_t rot_dim = pe->ne[2] * 2;
|
int64_t rot_dim = pe->ne[2] * 2;
|
||||||
GGML_ASSERT(rot_dim <= x->ne[0]);
|
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),
|
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
|
||||||
pe,
|
pe,
|
||||||
false);
|
false);
|
||||||
@@ -209,8 +210,8 @@ namespace MiniMaxH3 {
|
|||||||
q = q_norm->forward(ctx, q);
|
q = q_norm->forward(ctx, q);
|
||||||
k = k_norm->forward(ctx, k);
|
k = k_norm->forward(ctx, k);
|
||||||
if (pe != nullptr) {
|
if (pe != nullptr) {
|
||||||
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
|
q = apply_partial_rope(ctx, q, pe);
|
||||||
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
|
k = apply_partial_rope(ctx, k, pe);
|
||||||
} else {
|
} else {
|
||||||
q = attention_layout(ctx->ggml_ctx, q);
|
q = attention_layout(ctx->ggml_ctx, q);
|
||||||
k = attention_layout(ctx->ggml_ctx, k);
|
k = attention_layout(ctx->ggml_ctx, k);
|
||||||
|
|||||||
@@ -188,8 +188,8 @@ namespace Qwen {
|
|||||||
auto v = project("to_v");
|
auto v = project("to_v");
|
||||||
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
|
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
|
||||||
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
|
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
|
||||||
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
|
q = Rope::apply_rope(ctx, q, pe);
|
||||||
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
|
k = Rope::apply_rope(ctx, k, pe);
|
||||||
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
|
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
|
||||||
// Preserve query-first attention evaluation while writing each layer's
|
// Preserve query-first attention evaluation while writing each layer's
|
||||||
// prefix before its full-sequence K/V can accumulate across layers.
|
// prefix before its full-sequence K/V can accumulate across layers.
|
||||||
|
|||||||
@@ -228,11 +228,12 @@ namespace MiniMaxH3VAE {
|
|||||||
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
|
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* x,
|
||||||
ggml_tensor* pe) {
|
ggml_tensor* pe) {
|
||||||
|
auto ctx = runner_ctx->ggml_ctx;
|
||||||
int64_t rot_dim = pe->ne[2] * 2;
|
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),
|
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
|
||||||
pe,
|
pe,
|
||||||
false);
|
false);
|
||||||
@@ -289,8 +290,8 @@ namespace MiniMaxH3VAE {
|
|||||||
batch_size);
|
batch_size);
|
||||||
q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
|
q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
|
||||||
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
|
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
|
||||||
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
|
q = apply_partial_rope(ctx, q, pe);
|
||||||
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
|
k = apply_partial_rope(ctx, k, pe);
|
||||||
auto out = ggml_ext_attention_ext(ctx,
|
auto out = ggml_ext_attention_ext(ctx,
|
||||||
q,
|
q,
|
||||||
k,
|
k,
|
||||||
|
|||||||
Reference in new issue
Block a user