diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
index ca9e3ec6e110..89f9e2e04edb 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -2153,21 +2153,12 @@ def __init__(
             self.value_cache.append(torch.zeros(1, KV_H, max_cached_seq_len, V_D, device=device, dtype=dtype))
             self.batch_reserve(self.paged_attentions[i], torch.tensor([max_cache_len for _ in range(batch_size)]))
 
-        def generate_causal_offset(offset: torch.Tensor):
-            def causal_offset_mask(b, h, q_idx, kv_idx):
-                return (offset + q_idx) >= kv_idx
-
-            return causal_offset_mask
-
         self.batch_size = batch_size
         self.max_cache_len = max_cache_len
         self.block_masks = []
+        block_mask = create_block_mask(noop_mask, batch_size, 1, 1, max_cache_len, device=device, BLOCK_SIZE=page_size)
         for i in range(max_cache_len):
-            mod = generate_causal_offset(
-                torch.tensor(i, device=device, dtype=torch.int32)
-            )
-            block_mask = create_block_mask(mod, batch_size, 1, 1, max_cache_len, device=device, BLOCK_SIZE=page_size)
-            self.block_masks.append(self.paged_attentions[0].convert_logical_block_mask(block_mask))
+            self.block_masks.append(self.paged_attentions[0].convert_logical_block_mask(block_mask, kv_len=torch.tensor([i]*batch_size)))
         self.score_mods = []
         self.score_mods.append(None)
         self.score_mods.append(None)