diff --git a/examples/models/llama/runner/static_attention_io_manager.h b/examples/models/llama/runner/static_attention_io_manager.h
index b077f414f02..55543379ebb 100644
--- a/examples/models/llama/runner/static_attention_io_manager.h
+++ b/examples/models/llama/runner/static_attention_io_manager.h
@@ -759,7 +759,7 @@ class StaticAttentionIOManager {
   }
 
   void set_lookahead_decoding_mask(
-      StaticAttentionMask<MaskT>& mask,
+      StaticAttentionMask<MaskT, MaskAllocatorT>& mask,
       size_t input_len,
       size_t ngram_size,
       size_t window_size,