Internal changes

PiperOrigin-RevId: 837775282

gemma/flash_attention.cc

@@ -17,13 +17,16 @@
 #include <stdint.h>
 
 #include <algorithm>
+#include <array>
 #include <cmath>
+#include <cstdlib>
 #include <limits>
 
 #include "compression/types.h"  // GEMMA_DISABLED_TARGETS
 #include "gemma/flash_structs.h"
 #include "util/threading_context.h"
 #include "util/zones.h"
+#include "hwy/base.h"
 #ifndef HWY_DISABLED_TARGETS
 #define HWY_DISABLED_TARGETS GEMMA_DISABLED_TARGETS
 #endif  // HWY_DISABLED_TARGETS

gemma/flash_attention.h

@@ -37,7 +37,7 @@ namespace gcpp {
       const AttentionActivationsPtrs& activations, ThreadingContext& ctx);     \
                                                                                \
   void SingleFlashAttention(size_t start_pos, size_t last_pos,                 \
-                            const float* HWY_RESTRICT q,                      \
+                            const BF16* HWY_RESTRICT q,                        \
                             const MatPtrT<KV_t>& k, const MatPtrT<KV_t>& v,    \
                             size_t layer_idx,                                  \
                             const AttentionActivationsPtrs& activations,       \
@@ -60,6 +60,7 @@ namespace gcpp {
                       size_t layer_idx, const MatPtr& query_norm_scale,        \
                       AttentionActivationsPtrs& activations, QBatch& qbatch,   \
                       ThreadingContext& ctx);                                  \
+                                                                               \
   /* NOLINTNEXTLINE(google-readability-namespace-comments) */                  \
   }  // namespace NAMESPACE
 

ops/ops-inl.h

@@ -25,6 +25,7 @@
 #include <cstdint>
 #include <random>
 #include <type_traits>  // std::enable_if_t
+#include <utility>
 #include <vector>
 
 #include "ops/matmul.h"