Merge branch 'dev' into upgrade-github-actions-node24

Author: salmanmkc

BUILD.bazel

@@ -66,6 +66,7 @@ cc_library(
     srcs = ["util/topology.cc"],
     hdrs = ["util/topology.h"],
     deps = [
+        "@highway//:bit_set",
         "@highway//:hwy",
         "@highway//:topology",
     ],
@@ -111,6 +112,7 @@ cc_library(
         ":threading",
         ":topology",
         ":zones",
+        "//io",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
         "@highway//:profiler",
@@ -139,7 +141,7 @@ cc_test(
         ":kv_cache",
         ":mat",
         ":matmul",
-        ":ops",
+        ":test_util",
         ":threading_context",
         ":weights",
         "@googletest//:gtest_main",  # buildcleaner: keep
@@ -172,8 +174,11 @@ cc_library(
     name = "test_util",
     hdrs = ["util/test_util.h"],
     deps = [
+        ":basics",
+        ":mat",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
+        "@highway//:nanobenchmark",
         "@highway//:stats",
     ],
 )
@@ -440,9 +445,9 @@ cc_test(
         ":gemma_lib",
         ":mat",
         ":ops",
+        ":query",
         ":test_util",
         ":threading_context",
-        ":zones",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "//compression:test_util",
         "//compression:types",
@@ -519,32 +524,70 @@ cc_library(
     ],
 )
 
+cc_test(
+    name = "kv_cache_test",
+    srcs = ["gemma/kv_cache_test.cc"],
+    deps = [
+        ":configs",
+        ":gemma_args",
+        ":kv_cache",
+        ":threading_context",
+        "//testing/base/public:gunit_main",
+        "@highway//:hwy",
+    ],
+)
+
+cc_library(
+    name = "query",
+    hdrs = ["gemma/query.h"],
+    deps = [
+        ":basics",
+        ":gemma_args",
+        ":kv_cache",
+        "@highway//:hwy",
+    ],
+)
+
 cc_library(
     name = "gemma_args",
     hdrs = ["gemma/gemma_args.h"],
     deps = [
         ":args",
         ":basics",
+        ":configs",
         ":mat",
+        ":threading_context",
         "//io",
         "@highway//:hwy",
         "@highway//:profiler",
     ],
 )
 
+cc_test(
+    name = "gemma_args_test",
+    srcs = ["gemma/gemma_args_test.cc"],
+    deps = [
+        ":gemma_args",
+        "@googletest//:gtest_main",  # buildcleaner: keep
+    ],
+)
+
 cc_library(
     name = "gemma_lib",
     srcs = [
         "gemma/attention.cc",
         "gemma/flash_attention.cc",
         "gemma/gemma.cc",
+        "gemma/tensor_stats.cc",
         "gemma/vit.cc",
     ],
     hdrs = [
         "gemma/activations.h",
         "gemma/attention.h",
         "gemma/flash_attention.h",
+        "gemma/flash_structs.h",
         "gemma/gemma.h",
+        "gemma/tensor_stats.h",
         "gemma/vit.h",
     ],
     exec_properties = {
@@ -555,6 +598,7 @@ cc_library(
         "gemma/gemma-inl.h",
     ],
     deps = [
+        ":allocator",
         ":basics",
         ":configs",
         ":gemma_args",
@@ -564,6 +608,7 @@ cc_library(
         ":matmul_env",
         ":model_store",
         ":ops",
+        ":query",
         ":threading",
         ":threading_context",
         ":weights",
@@ -577,8 +622,34 @@ cc_library(
         "@highway//:hwy",
         "@highway//:nanobenchmark",  # timer
         "@highway//:profiler",
+        "@highway//:stats",
         "@highway//:thread_pool",
         "@highway//hwy/contrib/sort:vqsort",
+    ] +
+    [
+    ],
+)
+
+cc_test(
+    name = "gemma_lib_test",
+    srcs = ["gemma/attention_test.cc"],
+    # MatMulEnvs are up to 20GB large.
+    tags = ["requires-mem:28g"],
+    deps = [
+        ":configs",
+        ":gemma_args",
+        ":gemma_lib",
+        ":kv_cache",
+        ":mat",
+        ":matmul_env",
+        ":test_util",
+        ":threading_context",
+        ":weights",
+        "@googletest//:gtest_main",  # buildcleaner: keep
+        "//compression:compress",
+        "//compression:types",
+        "@highway//:hwy",
+        "@highway//:hwy_test_util",
     ],
 )
 
@@ -604,7 +675,6 @@ cc_library(
         ":gemma_args",
         ":gemma_lib",
         ":matmul_env",
-        ":ops",
         ":threading_context",
         ":tokenizer",
         "@google_benchmark//:benchmark",

CMakeLists.txt

@@ -22,7 +22,7 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 2a16a50ff61071bb25ddef0ce35d92b0e2b9c579 EXCLUDE_FROM_ALL)
+FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 3b680cde3a556bead9cc23c8f595d07a44d5a0d5 EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(highway)
 
 ## Note: absl needs to be installed by sentencepiece. This will only happen if
@@ -82,6 +82,7 @@ set(SOURCES
   gemma/configs.h
   gemma/flash_attention.cc
   gemma/flash_attention.h
+  gemma/flash_structs.h
   gemma/gemma_args.h
   gemma/gemma-inl.h
   gemma/gemma.cc
@@ -221,6 +222,7 @@ set(GEMMA_TEST_FILES
   compression/nuq_test.cc
   compression/sfp_test.cc
   evals/gemma_test.cc
+  gemma/gemma_args_test.cc
   gemma/flash_attention_test.cc
   gemma/tensor_info_test.cc
   io/blob_store_test.cc

MODULE.bazel

@@ -18,7 +18,7 @@ bazel_dep(name = "google_benchmark", version = "1.8.5")
 # Require a more recent version.
 git_override(
     module_name = "highway",
-    commit = "2a16a50ff61071bb25ddef0ce35d92b0e2b9c579",
+    commit = "3b680cde3a556bead9cc23c8f595d07a44d5a0d5",
     remote = "https://github.com/google/highway",
 )
 

README.md

@@ -55,7 +55,6 @@ Guidelines](https://opensource.google.com/conduct/).
 
     -   CPU-only inference for: Gemma 2-3, PaliGemma 2.
     -   Sampling with TopK and temperature.
-    -   Backward pass (VJP) and Adam optimizer for Gemma research.
 
 -   Optimizations
 
@@ -452,7 +451,7 @@ FetchContent_MakeAvailable(sentencepiece)
 FetchContent_Declare(gemma GIT_REPOSITORY https://github.com/google/gemma.cpp GIT_TAG origin/main)
 FetchContent_MakeAvailable(gemma)
 
-FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 2a16a50ff61071bb25ddef0ce35d92b0e2b9c579)
+FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 3b680cde3a556bead9cc23c8f595d07a44d5a0d5)
 FetchContent_MakeAvailable(highway)
 ```
 
@@ -520,13 +519,19 @@ Mikhaylov, Eugene Kliuchnikov, Jan Wassenberg, Jyrki Alakuijala, Lode
 Vandevenne, Luca Versari, Martin Bruse, Phil Culliton, Sami Boukortt, Thomas
 Fischbacher and Zoltan Szabadka. It was removed in 2025-09.
 
-Gemma-2 support was implemented in June/July 2024 with the help of several
-people.
+Gemma 2 support was implemented in June/July 2024 with the help of several
+people including Daniel Keysers and Phil Culliton.
 
 PaliGemma support was implemented in September 2024 with contributions from
 Daniel Keysers.
 
+Gemma 3 support was implemented in January-March 2025 with contributions from
+Daniel Keysers and Phil Culliton.
+
 [Jan Wassenberg](mailto:janwas@google.com) has continued to contribute many
 improvements, including major gains in efficiency, since the initial release.
 
+[Phil Culliton](mailto:philculliton@google.com) has worked on model releases,
+eval and validation, GTM, and quantization, since the initial release.
+
 This is not an officially supported Google product.

compression/BUILD.bazel

@@ -1,4 +1,4 @@
-# Weight compression and analysis.
+# Compressed tensor types.
 
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("@rules_cc//cc:cc_test.bzl", "cc_test")
@@ -101,13 +101,11 @@ cc_test(
     # for test_suite.
     tags = ["hwy_ops_test"],
     deps = [
-        ":distortion",
         ":int",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "//:test_util",
         "@highway//:hwy",
         "@highway//:hwy_test_util",
-        "@highway//:nanobenchmark",
     ],
 )
 
@@ -135,8 +133,8 @@ cc_test(
     # for test_suite.
     tags = ["hwy_ops_test"],
     deps = [
+        ":compress",
         ":distortion",
-        ":sfp",
         "@googletest//:gtest_main",  # buildcleaner: keep
         "//:test_util",
         "@highway//:hwy",
@@ -182,7 +180,6 @@ cc_library(
         "//:mat",
         "//:threading_context",
         "@highway//:hwy",
-        "@highway//:nanobenchmark",
         "@highway//:profiler",
         "@highway//:stats",
         "@highway//:thread_pool",
@@ -209,19 +206,3 @@ cc_test(
         "@highway//:hwy_test_util",
     ],
 )
-
-# For internal experimentation
-cc_library(
-    name = "analyze",
-    textual_hdrs = ["analyze.h"],
-    deps = [
-        ":int",
-        ":nuq",
-        ":sfp",
-        ":types",
-        "@highway//:hwy",
-        "@highway//:stats",
-        "@highway//:thread_pool",
-        "@highway//hwy/contrib/sort:vqsort",
-    ],
-)