mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 03:24:30 +00:00
perf: avoid redundant MOSS QKV and value layout copies (#2209)
Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
1 parent
2a5f81095d
commit
e6370bf093
2 files changed
+10
-7
No files matched your search
@@ -678,21 +678,23 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
|
||||
// CUDA FA MMA starts at 64 channels; keep the original head's attention scale.
|
||||
q_in = ggml_pad(ctx, q_in, static_cast<int>(64 - d_head), 0, 0, 0);
|
||||
k_in = ggml_pad(ctx, k_in, static_cast<int>(64 - d_head), 0, 0, 0);
|
||||
v_in = ggml_pad(ctx, v_in, static_cast<int>(64 - d_head), 0, 0, 0);
|
||||
}
|
||||
if (kv_scale != 1.0f) {
|
||||
k_in = ggml_ext_scale(ctx, k_in, kv_scale);
|
||||
}
|
||||
k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16);
|
||||
|
||||
v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
|
||||
v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
|
||||
if (pad_head) {
|
||||
v_in = ggml_pad(ctx, v_in, static_cast<int>(64 - d_head), 0, 0, 0);
|
||||
}
|
||||
if (kv_scale != 1.0f) {
|
||||
v_in = ggml_ext_scale(ctx, v_in, kv_scale);
|
||||
}
|
||||
v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
|
||||
v_in = ggml_permute(ctx, v_in, 0, 2, 1, 3);
|
||||
auto v_f16 = ggml_cast(ctx, v_in, GGML_TYPE_F16);
|
||||
if (!ggml_backend_supports_op(backend, v_f16)) {
|
||||
v_f16 = ggml_cast(ctx, ggml_ext_cont(ctx, v_in), GGML_TYPE_F16);
|
||||
}
|
||||
v_in = v_f16;
|
||||
v_in = ggml_reshape_3d(ctx, v_in, v_in->ne[0], L_k, n_kv_head * N);
|
||||
|
||||
if (mask_in != nullptr) {
|
||||
// ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped
|
||||
|
||||
Reference in new issue
Block a user