diff --git a/gemma/flash_attention.h b/gemma/flash_attention.h
index 47618c48..7e83a823 100644
--- a/gemma/flash_attention.h
+++ b/gemma/flash_attention.h
@@ -22,6 +22,7 @@
 
 #include <cstdint>
 
+#include "gemma/activations.h"
 #include "gemma/configs.h"
 #include "gemma/flash_structs.h"
 #include "gemma/kv_cache.h"