Internal changes

PiperOrigin-RevId: 853138600

BUILD.bazel

@@ -608,7 +608,9 @@ cc_library(
     ],
     deps = [
         ":activations",
+        ":basics",
         ":configs",
+        ":kv_cache",
         ":mat",
         ":matmul",
         ":matmul_env",

gemma/flash_attention.cc

@@ -22,10 +22,14 @@
 #include <cstdlib>
 #include <iostream>
 #include <limits>
+#include <type_traits>
 #include <vector>
 
 #include "compression/types.h"  // GEMMA_DISABLED_TARGETS
 #include "gemma/flash_structs.h"
+#include "gemma/kv_cache.h"
+#include "gemma/query.h"
+#include "util/basics.h"
 #include "util/threading_context.h"
 #include "util/zones.h"
 #include "hwy/base.h"

util/test_util.h

@@ -115,7 +115,7 @@ template <typename T>
 void PrintMatPtr(MatPtrT<T> mat) {
   for (int i = 0; i < mat.Rows(); ++i) {
     for (int j = 0; j < mat.Cols(); ++j) {
-      std::cerr << mat.Row(i)[j] << " ,";
+      std::cerr << hwy::ConvertScalarTo<float>(mat.Row(i)[j]) << " ,";
     }
     std::cerr << std::endl;
   }