Merge 6ce8e7028c into 1fe009caed
This commit is contained in:
commit
64b0f8f71e
|
|
@ -2074,6 +2074,18 @@ static struct ggml_cgraph * whisper_build_graph_encoder(
|
|||
|
||||
const float KQscale = 1.0f/sqrtf(float(n_state_head));
|
||||
|
||||
// with flash attention the self-attention K/V are viewed zero-padded to
|
||||
// n_ctx_pad, but only n_ctx rows are written - the padding holds stale
|
||||
// data from previous calls, so it must be masked out
|
||||
struct ggml_tensor * KQ_mask_f16 = nullptr;
|
||||
if (wctx.params.flash_attn) {
|
||||
struct ggml_tensor * KQ_mask = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_ctx_pad, n_ctx, 1);
|
||||
ggml_set_name(KQ_mask, "KQ_mask_enc");
|
||||
ggml_set_input(KQ_mask);
|
||||
|
||||
KQ_mask_f16 = ggml_cast(ctx0, KQ_mask, GGML_TYPE_F16);
|
||||
}
|
||||
|
||||
// ===================================================================
|
||||
// NOTE: experimenting with partial evaluation of the encoder (ignore)
|
||||
//static int iter = -1;
|
||||
|
|
@ -2162,7 +2174,7 @@ static struct ggml_cgraph * whisper_build_graph_encoder(
|
|||
ggml_element_size(kv_pad.v)*n_state_head,
|
||||
0);
|
||||
|
||||
cur = ggml_flash_attn_ext(ctx0, Q, K, V, nullptr, KQscale, 0.0f, 0.0f);
|
||||
cur = ggml_flash_attn_ext(ctx0, Q, K, V, KQ_mask_f16, KQscale, 0.0f, 0.0f);
|
||||
|
||||
cur = ggml_reshape_2d(ctx0, cur, n_state, n_ctx);
|
||||
} else {
|
||||
|
|
@ -2434,6 +2446,28 @@ static bool whisper_encode_internal(
|
|||
return false;
|
||||
}
|
||||
|
||||
// mask out the padded part of the flash-attention K/V views
|
||||
{
|
||||
struct ggml_tensor * KQ_mask = ggml_graph_get_tensor(gf, "KQ_mask_enc");
|
||||
if (KQ_mask) {
|
||||
const int n_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : wctx.model.hparams.n_audio_ctx;
|
||||
const int n_ctx_pad = GGML_PAD(n_ctx, 256);
|
||||
|
||||
wstate.inp_mask.resize(ggml_nelements(KQ_mask));
|
||||
|
||||
float * data = wstate.inp_mask.data();
|
||||
memset(data, 0, ggml_nbytes(KQ_mask));
|
||||
|
||||
for (int j = 0; j < n_ctx; ++j) {
|
||||
for (int i = n_ctx; i < n_ctx_pad; ++i) {
|
||||
data[j*n_ctx_pad + i] = -INFINITY;
|
||||
}
|
||||
}
|
||||
|
||||
ggml_backend_tensor_set(KQ_mask, wstate.inp_mask.data(), 0, ggml_nelements(KQ_mask)*sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
if (!ggml_graph_compute_helper(sched, gf, n_threads)) {
|
||||
return false;
|
||||
}
|
||||
|
|
@ -2517,6 +2551,18 @@ static struct ggml_cgraph * whisper_build_graph_decoder(
|
|||
|
||||
struct ggml_tensor * KQ_mask_f16 = ggml_cast(ctx0, KQ_mask, GGML_TYPE_F16);
|
||||
|
||||
// with flash attention the cross-attention K/V are viewed zero-padded to
|
||||
// n_audio_ctx_pad, but only n_audio_ctx rows are written - the padding
|
||||
// holds stale data from previous calls, so it must be masked out
|
||||
struct ggml_tensor * KQ_mask_cross_f16 = nullptr;
|
||||
if (wctx.params.flash_attn) {
|
||||
struct ggml_tensor * KQ_mask_cross = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_audio_ctx_pad, n_tokens, 1);
|
||||
ggml_set_name(KQ_mask_cross, "KQ_mask_cross");
|
||||
ggml_set_input(KQ_mask_cross);
|
||||
|
||||
KQ_mask_cross_f16 = ggml_cast(ctx0, KQ_mask_cross, GGML_TYPE_F16);
|
||||
}
|
||||
|
||||
// token encoding + position encoding
|
||||
struct ggml_tensor * cur =
|
||||
ggml_add(ctx0,
|
||||
|
|
@ -2698,7 +2744,7 @@ static struct ggml_cgraph * whisper_build_graph_decoder(
|
|||
ggml_element_size(wstate.kv_cross.v)*n_state_head,
|
||||
ggml_element_size(wstate.kv_cross.v)*n_state*n_audio_ctx_pad*il);
|
||||
|
||||
cur = ggml_flash_attn_ext(ctx0, Q, Kcross, Vcross, nullptr, KQscale, 0.0f, 0.0f);
|
||||
cur = ggml_flash_attn_ext(ctx0, Q, Kcross, Vcross, KQ_mask_cross_f16, KQscale, 0.0f, 0.0f);
|
||||
|
||||
cur = ggml_reshape_2d(ctx0, cur, n_state, n_tokens);
|
||||
} else {
|
||||
|
|
@ -2945,6 +2991,28 @@ static bool whisper_decode_internal(
|
|||
ggml_backend_tensor_set(KQ_mask, wstate.inp_mask.data(), 0, ggml_nelements(KQ_mask)*sizeof(float));
|
||||
}
|
||||
|
||||
// mask out the padded part of the flash-attention cross K/V views
|
||||
{
|
||||
struct ggml_tensor * KQ_mask_cross = ggml_graph_get_tensor(gf, "KQ_mask_cross");
|
||||
if (KQ_mask_cross) {
|
||||
const int n_audio_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : wctx.model.hparams.n_audio_ctx;
|
||||
const int n_audio_ctx_pad = GGML_PAD(n_audio_ctx, 256);
|
||||
|
||||
wstate.inp_mask.resize(ggml_nelements(KQ_mask_cross));
|
||||
|
||||
float * data = wstate.inp_mask.data();
|
||||
memset(data, 0, ggml_nbytes(KQ_mask_cross));
|
||||
|
||||
for (int j = 0; j < n_tokens; ++j) {
|
||||
for (int i = n_audio_ctx; i < n_audio_ctx_pad; ++i) {
|
||||
data[j*n_audio_ctx_pad + i] = -INFINITY;
|
||||
}
|
||||
}
|
||||
|
||||
ggml_backend_tensor_set(KQ_mask_cross, wstate.inp_mask.data(), 0, ggml_nelements(KQ_mask_cross)*sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
logits = ggml_graph_node(gf, -1);
|
||||
|
||||
if (!ggml_graph_compute_helper(sched, gf, n_threads)) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue