Minor cleanup, on-demand NUQ buffer allocation

threading_context: add profiler
compress-inl: add constexpr, on-demand alloc NUQ buffer
gemma_py: model->gemma
Move ScaleWeights to compress.cc
Move PromptWrapping to configs.h
PiperOrigin-RevId: 748277092

Author: jan-wassenberg

.github/workflows/build.yml

@@ -82,7 +82,7 @@ jobs:
         subprocess.run(["cp", "/kaggle/input/gemma-build-artifacts/gemma", "/kaggle/working"])
         subprocess.run(["chmod", "700", "/kaggle/working/gemma"])
         subprocess.run(["cp", "/kaggle/input/gemma-build-artifacts/_deps/sentencepiece-build/src/libsentencepiece.so.0", "/kaggle/working"])
-        output = subprocess.run(["/kaggle/working/gemma", "--tokenizer", "/kaggle/input/gemma/gemmacpp/2b-it-sfp/4/tokenizer.spm", "--compressed_weights", "/kaggle/input/gemma/gemmacpp/2b-it-sfp/4/2b-it-sfp.sbs", "--model", "2b-it", "--verbosity", "0", "--max_generated_tokens", "128"], stdout=subprocess.PIPE, input='Write an email to the moon.', encoding='ascii').stdout
+        output = subprocess.run(["/kaggle/working/gemma", "--tokenizer", "/kaggle/input/gemma/gemmacpp/2b-it-sfp/4/tokenizer.spm", "--weights", "/kaggle/input/gemma/gemmacpp/2b-it-sfp/4/2b-it-sfp.sbs", "--model", "2b-it", "--verbosity", "0", "--max_generated_tokens", "128"], stdout=subprocess.PIPE, input='Write an email to the moon.', encoding='ascii').stdout
         assert("write an email to the moon." not in output.lower());
         assert("moon" in output.lower());
         EOF

BUILD.bazel

@@ -92,6 +92,8 @@ cc_library(
         ":basics",
         ":threading",
         ":topology",
+        "@highway//:hwy",
+        "@highway//:profiler",
     ],
 )
 
@@ -180,6 +182,7 @@ cc_library(
         "//compression:shared",
         "@highway//:hwy",
         "@highway//:profiler",
+        "@highway//:thread_pool",
     ],
 )
 
@@ -664,6 +667,7 @@ cc_test(
         ":mat",
         ":prompt",
         ":sampler",
+        ":threading_context",
         ":weights",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "@highway//:thread_pool",

compression/BUILD.bazel

@@ -70,6 +70,7 @@ cc_library(
     hdrs = ["blob_store.h"],
     deps = [
         ":io",
+        "//:basics",
         "//:threading_context",
         "@highway//:hwy",
         "@highway//:thread_pool",
@@ -130,7 +131,6 @@ cc_library(
     textual_hdrs = ["sfp-inl.h"],
     deps = [
         ":shared",
-        "//:basics",
         "@highway//:hwy",
     ],
 )
@@ -195,7 +195,6 @@ cc_test(
     deps = [
         ":distortion",
         ":nuq",
-        ":sfp",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "//:test_util",
         "@highway//:hwy",
@@ -225,6 +224,7 @@ cc_library(
         "//:mat",
         "@highway//:hwy",
         "@highway//:nanobenchmark",
+        "@highway//:profiler",
         "@highway//:stats",
         "@highway//:thread_pool",
     ],
@@ -259,6 +259,7 @@ cc_library(
     deps = [
         ":nuq",
         ":sfp",
+        ":shared",
         "@highway//:hwy",
         "@highway//:stats",
         "@highway//:thread_pool",

compression/compress-inl.h

@@ -21,8 +21,7 @@
 #include <stdint.h>
 #include <stdio.h>
 
-#include <cmath>  // lroundf, only if COMPRESS_STATS
-#include <string>
+#include <memory>
 #include <vector>
 
 #include "compression/blob_store.h"
@@ -35,6 +34,10 @@
 #include "hwy/contrib/thread_pool/thread_pool.h"
 #include "hwy/timer.h"
 
+#if COMPRESS_STATS
+#include <cmath>  // lroundf
+#endif
+
 #endif  // THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_INL_H_
 
 // Include guard for (potentially) SIMD code.
@@ -388,7 +391,7 @@ struct CompressTraits<SfpStream> {
                                   const size_t packed_ofs) {
     SfpCodec::Enc(df, raw, num, packed.ptr + packed_ofs);
 
-    if (COMPRESS_STATS) {
+    if constexpr (COMPRESS_STATS) {
       const hn::Repartition<BF16, DF> dbf;
       auto distorted =
           hwy::AllocateAligned<BF16>(hwy::RoundUpTo(num, hn::Lanes(dbf)));
@@ -432,9 +435,10 @@ struct CompressTraits<NuqStream> {
                                   size_t num, CompressPerThread& tls,
                                   const PackedSpan<Packed>& packed,
                                   const size_t packed_ofs) {
-    NuqCodec::Enc(df, raw, num, tls.buf, packed, packed_ofs);
+    if (!tls.buf) tls.buf = std::make_unique<NuqStream::ClusterBuf>();
+    NuqCodec::Enc(df, raw, num, *tls.buf, packed, packed_ofs);
 
-    if (COMPRESS_STATS) {
+    if constexpr (COMPRESS_STATS) {
       for (size_t i = 0; i < num; ++i) {
         tls.stats.NotifyIn(static_cast<int>(lroundf(raw[i] * 100.0f + 500.0f)));
       }
@@ -478,7 +482,7 @@ HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
                            const size_t packed_ofs, hwy::ThreadPool& pool) {
   packed.BoundsCheck(packed_ofs, num);
   work.tls.resize(pool.NumWorkers());
-  if (COMPRESS_STATS) {
+  if constexpr (COMPRESS_STATS) {
     for (auto& tls : work.tls) {
       tls.stats.Reset();
     }
@@ -487,7 +491,7 @@ HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
   const bool want_bench = COMPRESS_STATS || !kIsTest;
   const double t0 = want_bench ? hwy::platform::Now() : 0.0;
 
-  using Traits = CompressTraits<Packed>;
+  using Traits = CompressTraits<hwy::RemoveConst<Packed>>;
   constexpr size_t kBatch = 8192;
   const size_t num_batches = hwy::DivCeil(num, kBatch);
   pool.Run(0, num_batches,
@@ -508,7 +512,7 @@ HWY_NOINLINE void Compress(const float* HWY_RESTRICT raw, size_t num,
     fprintf(stderr, "Compress %.1f MB/s\n", mbps);
   }
 
-  if (COMPRESS_STATS) {
+  if constexpr (COMPRESS_STATS) {
     for (size_t i = 1; i < work.tls.size(); ++i) {
       work.tls[0].stats.Assimilate(work.tls[i].stats);
     }
@@ -534,7 +538,7 @@ void Compress2(DF df, VF raw0, VF raw1, const PackedSpan<Packed>& packed,
                const size_t packed_ofs) {
   static_assert(hwy::IsSameEither<Packed, float, BF16>());
   packed.BoundsCheck(packed_ofs, 2 * hn::Lanes(df));
-  using Traits = CompressTraits<Packed>;
+  using Traits = CompressTraits<hwy::RemoveConst<Packed>>;
   Traits::Store2(df, raw0, raw1, packed, packed_ofs);
 }
 

compression/compress.cc

@@ -15,8 +15,34 @@
 
 #include "compression/compress.h"
 
+#include <stddef.h>
+#include <stdint.h>
+
+#include "util/mat.h"
+#include "hwy/base.h"
+#include "hwy/profiler.h"
+
 namespace gcpp {
 
-// TODO: move ScaleWeights here.
+float ScaleWeights(float* HWY_RESTRICT raw, size_t num) {
+  PROFILER_FUNC;
+
+  float maxabs = 0.0;
+  for (size_t i = 0; i < num; ++i) {
+    maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i]));
+  }
+  if (maxabs <= SfpStream::kMax) {
+    return 1.0f;
+  }
+  const float scale = maxabs / SfpStream::kMax;
+  const float inv_scale = static_cast<float>(1.0 / static_cast<double>(scale));
+  for (size_t i = 0; i < num; ++i) {
+    // Clamp because kMax may still be exceeded.
+    const float magn =
+        HWY_MIN(SfpStream::kMax, hwy::ScalarAbs(raw[i] * inv_scale));
+    raw[i] = hwy::ScalarCopySign(magn, raw[i]);
+  }
+  return scale;
+}
 
 }  // namespace gcpp

compression/compress.h

@@ -17,26 +17,19 @@
 #ifndef THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_H_
 #define THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_H_
 
-#include "hwy/base.h"
 #define COMPRESS_STATS 0
 
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
 
-#include <cstdio>
-#include <cstring>
-#include <string>
-#include <unordered_map>
-#include <utility>
+#include <memory>
 #include <vector>
 
-// IWYU pragma: begin_exports
 #include "compression/blob_store.h"
 #include "compression/fields.h"
 #include "compression/io.h"
-#include "compression/shared.h"
-#include "gemma/tensor_index.h"
+#include "compression/shared.h"  // NuqStream::ClusterBuf
 #include "util/basics.h"
 // IWYU pragma: end_exports
 #include "gemma/configs.h"
@@ -174,7 +167,8 @@ struct CompressStats {
 #endif  // COMPRESS_STATS
 
 struct CompressPerThread {
-  NuqStream::ClusterBuf buf;
+  // Allocated the first time NUQ is used.
+  std::unique_ptr<NuqStream::ClusterBuf> buf;
   CompressStats stats;
 };
 
@@ -375,5 +369,11 @@ class ReadFromBlobStore {
   std::vector<std::string> file_keys_;
 };
 
+// Returns 1.0f if all magnitudes are <= `SfpStream::kMax`, otherwise scales
+// them such that the largest magnitude is `SfpStream::kMax`, and returns the
+// multiplier with which to restore the original values. This is only necessary
+// before compressing to `SfpStream` and `NuqStream`.
+float ScaleWeights(float* HWY_RESTRICT raw, size_t num);
+
 }  // namespace gcpp
 #endif  // THIRD_PARTY_GEMMA_CPP_COMPRESSION_COMPRESS_H_

compression/shared.h

@@ -13,8 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Definitions shared between the public compress-inl.h interface and the
-// sfp-inl.h and nuq-inl.h implementation details.
+// Types shared between tensor definitions and `compress-inl.h`.
 
 #ifndef THIRD_PARTY_GEMMA_CPP_COMPRESSION_SHARED_H_
 #define THIRD_PARTY_GEMMA_CPP_COMPRESSION_SHARED_H_
@@ -63,30 +62,6 @@ struct SfpStream {
 };
 #pragma pack(pop)
 
-// Returns 1.0f if all magnitudes are <= SfpStream::kMax, otherwise scales them
-// such that the largest magnitude is SfpStream::kMax, and returns the
-// multiplier with which to restore the original values. This is only necessary
-// before compressing to SfpStream.
-// TODO: vectorize
-static inline float ScaleWeights(float* HWY_RESTRICT raw, size_t num) {
-  float maxabs = 0.0;
-  for (size_t i = 0; i < num; ++i) {
-    maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i]));
-  }
-  if (maxabs <= SfpStream::kMax) {
-    return 1.0f;
-  }
-  const float scale = maxabs / SfpStream::kMax;
-  const float inv_scale = static_cast<float>(1.0 / static_cast<double>(scale));
-  for (size_t i = 0; i < num; ++i) {
-    // Clamp because kMax may still be exceeded.
-    const float magn =
-        HWY_MIN(SfpStream::kMax, hwy::ScalarAbs(raw[i] * inv_scale));
-    raw[i] = hwy::ScalarCopySign(magn, raw[i]);
-  }
-  return scale;
-}
-
 // Non-uniform quantization: a compressed representation of f32 inputs that
 // supports seeking at a granularity of 1 (for `DecompressAndZeroPad`) or
 // two vectors (for `Decompress2`), and decoding to bf16/f32.
@@ -185,20 +160,6 @@ constexpr bool IsNuqStream() {
   return hwy::IsSame<hwy::RemoveCvRef<Packed>, NuqStream>();
 }
 
-// Instruction-tuned models require extra 'turn structure' tokens in prompts.
-enum class PromptWrapping {
-  GEMMA_IT,
-  GEMMA_PT,
-  GEMMA_VLM,
-  PALIGEMMA,
-  kSentinel  // must be last
-};
-
-inline bool EnumValid(PromptWrapping type) {
-  return static_cast<int>(type) >= 0 &&
-         static_cast<int>(type) < static_cast<int>(PromptWrapping::kSentinel);
-}
-
 // Tensor types for loading weights. Note that not all types are supported as
 // weights for a model, but can be used for other purposes, such as types for
 // `WeightsPtrs`. When adding a new type that is supported, also

gemma/configs.h

@@ -49,6 +49,20 @@ static constexpr size_t kMaxConv1DWidth = 4;
 
 using EmbedderInputT = BF16;
 
+// Instruction-tuned models require extra 'turn structure' tokens in prompts.
+enum class PromptWrapping {
+  GEMMA_IT,
+  GEMMA_PT,
+  GEMMA_VLM,
+  PALIGEMMA,
+  kSentinel  // must be last
+};
+
+static inline bool EnumValid(PromptWrapping wrapping) {
+  return static_cast<size_t>(wrapping) <
+         static_cast<size_t>(PromptWrapping::kSentinel);
+}
+
 enum class LayerAttentionType {
   kGemma,
   kGriffinRecurrentBlock,

ops/dot-inl.h

@@ -15,9 +15,6 @@
 
 #include <stddef.h>
 
-#include "compression/compress.h"
-#include "util/mat.h"
-#include "hwy/base.h"
 #include "hwy/profiler.h"
 
 // Include guard for (potentially) SIMD code.

paligemma/BUILD.bazel

@@ -40,6 +40,7 @@ cc_test(
     ],
     deps = [
         "@googletest//:gtest_main",  # buildcleaner: keep
+        "//:allocator",
         "//:benchmark_helper",
         "//:common",
         "//:gemma_lib",

paligemma/paligemma_test.cc

@@ -20,7 +20,9 @@
 #include "compression/shared.h"
 #include "evals/benchmark_helper.h"
 #include "gemma/common.h"
+#include "gemma/configs.h"
 #include "gemma/gemma.h"
+#include "util/allocator.h"
 #include "hwy/base.h"
 #include "hwy/tests/hwy_gtest.h"
 
@@ -50,17 +52,18 @@ class PaliGemmaTest : public ::testing::Test {
 
 void PaliGemmaTest::InitVit(const std::string& path) {
   ASSERT_NE(s_env->GetGemma(), nullptr);
-  Gemma& model = *(s_env->GetGemma());
-  image_tokens_ =
-      ImageTokens(Extents2D(model.GetModelConfig().vit_config.seq_len,
-                            model.GetModelConfig().model_dim));
+  const Allocator2& allocator = s_env->Env().ctx.allocator;
+  Gemma& gemma = *(s_env->GetGemma());
+  image_tokens_ = ImageTokens(
+      allocator, Extents2D(gemma.GetModelConfig().vit_config.seq_len,
+                           gemma.GetModelConfig().model_dim));
   Image image;
-  HWY_ASSERT(model.Info().wrapping == PromptWrapping::PALIGEMMA);
+  HWY_ASSERT(gemma.GetModelConfig().wrapping == PromptWrapping::PALIGEMMA);
   HWY_ASSERT(image.ReadPPM(path));
-  const size_t image_size = model.GetModelConfig().vit_config.image_size;
+  const size_t image_size = gemma.GetModelConfig().vit_config.image_size;
   image.Resize(image_size, image_size);
   RuntimeConfig runtime_config = {.gen = &s_env->MutableGen(), .verbosity = 0};
-  model.GenerateImageTokens(runtime_config, image, image_tokens_);
+  gemma.GenerateImageTokens(runtime_config, image, image_tokens_);
 }
 
 std::string PaliGemmaTest::GemmaReply(const std::string& prompt_text) const{
@@ -124,7 +127,7 @@ TEST_F(PaliGemmaTest, General) {
   };
   const char* (*qa)[2];
   size_t num;
-  switch (s_env->GetGemma()->Info().model) {
+  switch (s_env->GetGemma()->GetModelConfig().model) {
     case Model::PALIGEMMA_224:
       qa = kQA_3B_mix_224;
       num = sizeof(kQA_3B_mix_224) / sizeof(kQA_3B_mix_224[0]);

python/BUILD.bazel

@@ -21,10 +21,10 @@ pybind_extension(
     name = "gemma",
     srcs = ["gemma_py.cc"],
     deps = [
-        "//:allocator",
         "//:benchmark_helper",
         "//:gemma_args",
         "//:gemma_lib",
+        "//:threading_context",
         "//compression:shared",
         "@highway//:hwy",
     ],