diff --git a/server/text_generation_server/models/custom_modeling/flash_santacoder_modeling.py b/server/text_generation_server/models/custom_modeling/flash_santacoder_modeling.py
index 3a61d058..c0d0bbb3 100644
--- a/server/text_generation_server/models/custom_modeling/flash_santacoder_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_santacoder_modeling.py
@@ -363,6 +363,7 @@ class Block(nn.Module):
         slots,
         input_lengths,
         max_s,
+        prefill_cache_indices
     ):
         hidden_states, residual = self.ln_1(hidden_states, residual)
         hidden_states = self.attn(