diff --git a/src/layers/flash_attention.cc b/src/layers/flash_attention.cc
index 45b3a49ed..26881f5fa 100644
--- a/src/layers/flash_attention.cc
+++ b/src/layers/flash_attention.cc
@@ -18,6 +18,7 @@ namespace ctranslate2 {
     void FlashMultiHeadAttention::operator()(const StorageView& queries,
                                              const StorageView&,
                                              const StorageView* values_lengths,
+                                             const StorageView* values_offsets,
                                              StorageView& output,
                                              StorageView* cached_keys,
                                              StorageView* cached_values,
diff --git a/src/layers/wav2vec2bert.cc b/src/layers/wav2vec2bert.cc
index 7ac3620e5..3b5d7c87d 100644
--- a/src/layers/wav2vec2bert.cc
+++ b/src/layers/wav2vec2bert.cc
@@ -51,6 +51,7 @@ namespace ctranslate2 {
       (*_self_attention)(buffer1,
                          buffer1,
                          nullptr,
+                         nullptr,
                          buffer2,
                          nullptr,
                          nullptr,
@@ -147,6 +148,7 @@ namespace ctranslate2 {
       (*_self_attention)(buffer2,
                          buffer2,
                          nullptr,
+                         nullptr,
                          buffer1,
                          nullptr,
                          nullptr,