From b94c46f4352f839de738c2ff74e589dee92e269f Mon Sep 17 00:00:00 2001 From: Nils Reck Date: Fri, 14 Aug 2026 13:22:50 +0200 Subject: [PATCH] Make flash_attention effective for the Whisper encoder The flash_attention model option currently has no effect for Whisper: WhisperEncoderLayer constructs its TransformerEncoderLayer without forwarding model.use_flash_attention(), so the parameter always falls back to its default (false) and the encoder never uses the flash path. Forwarding the flag exposes a second issue: FlashMultiHeadAttention constructs ops::FlashAttention without is_causal, which defaults to true. That is correct for decoder self-attention (the only current user of the flash path) but wrong for any encoder, which must attend bidirectionally. Pass is_causal=_is_decoder instead; this is behavior-neutral for decoders and makes the flash path usable for encoder self-attention. Validated with a WITH_FLASH_ATTN=ON source build on one RTX PRO 6000 (Blackwell): Whisper encoder forward pass 358 ms -> 116 ms (3.1x), cosine similarity 0.999667 against the non-flash encoder output, pairwise WER -0.026 on our internal test set. --- src/layers/flash_attention.cc | 4 +++- src/layers/whisper.cc | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/layers/flash_attention.cc b/src/layers/flash_attention.cc index fb13b90fc..05812456e 100644 --- a/src/layers/flash_attention.cc +++ b/src/layers/flash_attention.cc @@ -110,7 +110,9 @@ namespace ctranslate2 { // init output StorageView context(dtype, device); - ops::FlashAttention fl_attn_ops(_queries_scale, _sliding_window); + // Causal masking only applies to decoder self-attention; encoder + // self-attention must attend bidirectionally. + ops::FlashAttention fl_attn_ops(_queries_scale, _sliding_window, /*is_causal=*/_is_decoder); fl_attn_ops(queries_proj, keys_proj, values_proj, context, cached_keys, cached_values, attention, return_normalized_attention, rotary_cos, rotary_sin, rotary_interleaved, nullptr/*alibli*/, offset); diff --git a/src/layers/whisper.cc b/src/layers/whisper.cc index 401c14677..763b47b28 100644 --- a/src/layers/whisper.cc +++ b/src/layers/whisper.cc @@ -17,7 +17,8 @@ namespace ctranslate2 { scope + "/layer", _num_heads, /*pre_norm=*/true, - ops::ActivationType::GELU)) + ops::ActivationType::GELU, + model.use_flash_attention())) , _output_norm(model, scope + "/layer_norm") { }