public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

whisper : mask padded K/V in flash-attention encoder and cross-attention 

With flash attention enabled, the encoder self-attention and the decoder
cross-attention view their K/V buffers padded to GGML_PAD(n_ctx, 256)
but only write n_ctx rows, and call ggml_flash_attn_ext with a null
mask. The unwritten pad rows are attended at full weight; their content
is whatever the buffer holds - zeros on a fresh state, or leftovers of
previous whisper_full calls with a different audio_ctx (layer strides
il*n_ctx_pad shift between window sizes, so runs overwrite each other's
pad regions).

On a long-running server this makes transcriptions depend on request
history: e.g. large-v3-turbo with audio_ctx=250 on a 3 s clip produced
5x sentence duplication on a fresh state, 2x after one audio_ctx=500
request, and clean output after full-window requests. The default full
window is affected too (36 unmasked rows out of 1536), just diluted
enough to go unnoticed.

Fix: give both call sites a proper mask input (mirroring the decoder
self-attention KQ_mask) with -INFINITY over [n_ctx, n_ctx_pad), filled
in whisper_encode_internal / whisper_decode_internal. Outputs at reduced
audio_ctx are now deterministic regardless of prior requests; timing is
unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Guenter Bartsch 2026-07-16 11:35:40 +02:00
parent 080bbbe852
commit 6ce8e7028c
1 changed files with 70 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
src/whisper.cpp
View File

@ -2068,6 +2068,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;
@ -2156,7 +2168,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 {
@ -2428,6 +2440,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;
}
@ -2511,6 +2545,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,
@ -2692,7 +2738,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 {
@ -2939,6 +2985,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)) {