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

+1 -1
Submodule ggml updated: 89c4413f5d...66c1a28760.
+15 -3
View File
@@ -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;
+2 -2
View File
@@ -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,
+6 -5
View File
@@ -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,
+5 -4
View File
@@ -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);
+2 -2
View File
@@ -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.
+5 -4
View File
@@ -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,