Prep for sharding gemma.cc: split into kv_cache, tokenizer.

jan-wassenberg · copybara-github · commit 09a7e75ead6f · 2024-07-02T09:31:06.000-07:00
Move activations.h to backprop/ to make space for another activations.h.

PiperOrigin-RevId: 648744500
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -94,13 +94,35 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "tokenizer",
+    srcs = ["gemma/tokenizer.cc"],
+    hdrs = ["gemma/tokenizer.h"],
+    deps = [
+        "//compression:io",
+        "@hwy//:hwy",
+        "@hwy//:nanobenchmark",  # timer
+        "@hwy//:profiler",
+        "@com_google_sentencepiece//:sentencepiece_processor",
+    ],
+)
+
+cc_library(
+    name = "kv_cache",
+    srcs = ["gemma/kv_cache.cc"],
+    hdrs = ["gemma/kv_cache.h"],
+    deps = [
+        ":common",
+        "@hwy//:hwy",
+    ],
+)
+
 cc_library(
     name = "gemma_lib",
     srcs = [
         "gemma/gemma.cc",
     ],
     hdrs = [
-        "gemma/activations.h",
         "gemma/gemma.h",
     ],
     exec_properties = {
@@ -114,6 +136,8 @@ cc_library(
     deps = [
         ":common",
         ":ops",
+        ":tokenizer",
+        ":kv_cache",
         ":weights",
         "//compression:compress",
         "//compression:io",
@@ -122,7 +146,6 @@ cc_library(
         "@hwy//:nanobenchmark",  # timer
         "@hwy//:profiler",
         "@hwy//:thread_pool",
-        "@com_google_sentencepiece//:sentencepiece_processor",
     ],
 )
 
@@ -321,6 +344,7 @@ cc_library(
         "backprop/forward.cc",
     ],
     hdrs = [
+        "backprop/activations.h",
         "backprop/backward.h",
         "backprop/backward-inl.h",
         "backprop/forward.h",
@@ -340,6 +364,7 @@ cc_library(
 cc_library(
     name = "backprop_scalar",
     hdrs = [
+        "backprop/activations.h",
         "backprop/backward_scalar.h",
         "backprop/common_scalar.h",
         "backprop/forward_scalar.h",
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -49,6 +49,7 @@ set(SOURCES
   compression/sfp.h
   compression/sfp-inl.h
   compression/test_util.h
+  backprop/activations.h
   backprop/backward.cc
   backprop/backward.h
   backprop/backward-inl.h
@@ -62,18 +63,21 @@ set(SOURCES
   backprop/optimizer.h
   evals/cross_entropy.cc
   evals/cross_entropy.h
-  gemma/configs.h
-  gemma/activations.h
   gemma/benchmark_helper.cc
   gemma/benchmark_helper.h
   gemma/common.cc
   gemma/common.h
+  gemma/configs.h
   gemma/gemma.cc
   gemma/gemma.h
+  gemma/kv_cache.cc
+  gemma/kv_cache.h
   gemma/ops.h
+  gemma/tokenizer.cc
+  gemma/tokenizer.h
+  gemma/weights_raw.h
   gemma/weights.cc
   gemma/weights.h
-  gemma/weights_raw.h
   util/app.h
   util/args.h
   )
diff --git a/backprop/activations.h b/backprop/activations.h
@@ -13,8 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_
-#define THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_
+#ifndef THIRD_PARTY_GEMMA_CPP_BACKPROP_ACTIVATIONS_H_
+#define THIRD_PARTY_GEMMA_CPP_BACKPROP_ACTIVATIONS_H_
 
 #include <stddef.h>
 
@@ -86,4 +86,4 @@ class ActivationsWrapper {
 
 }  // namespace gcpp
 
-#endif  // THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_
+#endif  // THIRD_PARTY_GEMMA_CPP_BACKPROP_ACTIVATIONS_H_
diff --git a/backprop/backward-inl.h b/backprop/backward-inl.h
@@ -22,11 +22,11 @@
 
 #include <stddef.h>
 
-#include <algorithm>
 #include <cmath>
+#include <vector>
 
+#include "backprop/activations.h"
 #include "backprop/prompt.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"
 #include "hwy/base.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
diff --git a/backprop/backward.cc b/backprop/backward.cc
@@ -15,8 +15,8 @@
 
 #include "backprop/backward.h"
 
+#include "backprop/activations.h"
 #include "backprop/prompt.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 
diff --git a/backprop/backward_scalar.h b/backprop/backward_scalar.h
@@ -22,9 +22,9 @@
 #include <cmath>
 #include <vector>
 
+#include "backprop/activations.h"
 #include "backprop/common_scalar.h"
 #include "backprop/prompt.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"  // EmbeddingScaling
 #include "gemma/weights_raw.h"
 
diff --git a/backprop/backward_scalar_test.cc b/backprop/backward_scalar_test.cc
@@ -26,12 +26,12 @@
 #include <vector>
 
 #include "gtest/gtest.h"
+#include "backprop/activations.h"
 #include "backprop/common_scalar.h"
 #include "backprop/forward_scalar.h"
 #include "backprop/prompt.h"
 #include "backprop/sampler.h"
 #include "backprop/test_util.h"
-#include "gemma/activations.h"
 #include "gemma/configs.h"
 #include "gemma/weights_raw.h"
 
diff --git a/backprop/backward_test.cc b/backprop/backward_test.cc
@@ -24,12 +24,12 @@
 #include <random>
 #include <vector>
 
+#include "backprop/activations.h"
 #include "backprop/backward_scalar.h"
 #include "backprop/common_scalar.h"
 #include "backprop/forward_scalar.h"
 #include "backprop/sampler.h"
 #include "backprop/test_util.h"
-#include "gemma/activations.h"
 #include "gemma/configs.h"
 #include "gemma/weights_raw.h"
 #include "hwy/base.h"
diff --git a/backprop/forward-inl.h b/backprop/forward-inl.h
@@ -23,7 +23,7 @@
 #include <cmath>
 #include <vector>
 
-#include "gemma/activations.h"
+#include "backprop/activations.h"
 #include "gemma/common.h"
 #include "gemma/configs.h"
 #include "hwy/base.h"
diff --git a/backprop/forward.cc b/backprop/forward.cc
@@ -15,8 +15,8 @@
 
 #include "backprop/forward.h"
 
+#include "backprop/activations.h"
 #include "backprop/prompt.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 
diff --git a/backprop/forward_scalar.h b/backprop/forward_scalar.h
@@ -23,9 +23,9 @@
 #include <complex>
 #include <vector>
 
+#include "backprop/activations.h"
 #include "backprop/common_scalar.h"
 #include "backprop/prompt.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"  // EmbeddingScaling
 #include "gemma/weights_raw.h"
 
diff --git a/backprop/optimize_test.cc b/backprop/optimize_test.cc
@@ -20,12 +20,12 @@
 #include <vector>
 
 #include "gtest/gtest.h"
+#include "backprop/activations.h"
 #include "backprop/backward.h"
 #include "backprop/forward.h"
 #include "backprop/optimizer.h"
 #include "backprop/prompt.h"
 #include "backprop/sampler.h"
-#include "gemma/activations.h"
 #include "gemma/common.h"
 #include "gemma/gemma.h"
 #include "gemma/weights.h"
diff --git a/gemma/common.h b/gemma/common.h
@@ -17,6 +17,7 @@
 #define THIRD_PARTY_GEMMA_CPP_GEMMA_COMMON_H_
 
 #include <math.h>  // sqrtf
+#include <stddef.h>
 #include <stdint.h>
 
 #include <string>
@@ -35,6 +36,12 @@ ByteStorageT AllocateSizeof() {
   return hwy::AllocateAligned<uint8_t>(sizeof(T));
 }
 
+constexpr size_t kPrefillBatchSize = 16;
+constexpr size_t kDecodeBatchSize = 1;
+constexpr size_t kBatchedQueryBatchSize = 16;
+constexpr size_t kMinAdjustedPrefillBatchSize =
+    HWY_MAX((size_t)1, kPrefillBatchSize / kBatchedQueryBatchSize);
+
 // Model variants: see configs.h for details.
 enum class Model {
   GEMMA_2B,
@@ -51,6 +58,13 @@ enum class ModelTraining { GEMMA_IT, GEMMA_PT };
 // Tensor types for loading weights.
 enum class Type { kF32, kBF16, kSFP };
 
+// TODO(janwas): merge with parser/ToString.
+struct ModelInfo {
+  Model model;
+  ModelTraining training;
+  Type weight;
+};
+
 // Returns the return value of FuncT<Config*<TWeight>>().operator()(args), where
 // Config* is selected via `model`. Typically called by CallForModelAndWeight,
 // but can also be called directly when FuncT does not actually use TWeight.
diff --git a/gemma/gemma.cc b/gemma/gemma.cc
diff --git a/gemma/gemma.h b/gemma/gemma.h
diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc
diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h
diff --git a/gemma/tokenizer.cc b/gemma/tokenizer.cc
diff --git a/gemma/tokenizer.h b/gemma/tokenizer.h
