Internal change.

BUILD.bazel

@@ -1067,7 +1067,8 @@ cc_test(
         ":configs",
         ":cross_entropy",
         ":gemma_lib",
-        "@googletest//:gtest_main",  # buildcleaner: keep
+        ":test_util",
+        "//testing/base/public:gunit_for_library_testonly",  # buildcleaner: keep
         "@highway//:hwy",
         "@highway//:hwy_test_util",
     ],
@@ -1087,7 +1088,7 @@ cc_test(
         ":benchmark_helper",
         ":gemma_lib",
         ":test_util",
-        "@googletest//:gtest_main",  # buildcleaner: keep
+        "//testing/base/public:gunit_for_library_testonly",  # buildcleaner: keep
         "@highway//:hwy",
         "@highway//:hwy_test_util",
         "@highway//:nanobenchmark",

compression/compress.cc

@@ -15,6 +15,7 @@
 
 #include "compression/compress.h"
 
+#include <cmath>
 #include <stddef.h>
 #include <stdint.h>
 

compression/python/compression_clif_aux.cc

@@ -85,8 +85,10 @@ class SbsWriterImpl : public ISbsWriter {
     }
 
     HWY_ASSERT(weights.size() == mat.Extents().Area());
-    Compress(weights.data(), weights.size(), working_set_, mat.Span(),
-             /*packed_ofs=*/0, ctx_);
+    {
+      Compress(weights.data(), weights.size(), working_set_, mat.Span(),
+               /*packed_ofs=*/0, ctx_);
+    }
     writer_.Add(name, mat.Packed(), mat.PackedBytes());
   }
 

compression/types.h

@@ -269,7 +269,8 @@ constexpr bool IsMxFp4Stream() {
 template <typename Packed>
 constexpr bool IsPacked() {
   return IsNuqStream<Packed>() || IsI8Stream<Packed>() ||
-         IsQ4_0Stream<Packed>() || IsMxFp4Stream<Packed>();
+         IsQ4_0Stream<Packed>() || IsMxFp4Stream<Packed>()
+         ;
 }
 
 template <typename Packed>
@@ -294,12 +295,14 @@ enum class Type {
   kInt8,
   kQ4_0,
   kMXFP4,
+  kReserved14,
 };
 // These are used in `ModelConfig.Specifier`, hence the strings will not
 // change, though new ones may be added.
 static constexpr const char* kTypeStrings[] = {
-    "unknown", "f32", "bf16", "sfp", "nuq",  "f64",  "u32",
-    "u64",     "i8",  "u16",  "u8",  "int8", "q4_0", "mxfp4"};
+    "unknown", "f32", "bf16", "sfp",   "nuq",   "f64",  "u32", "u64",
+    "i8",      "u16", "u8",   "int8",  "q4_0",  "mxfp4", "reserved14"
+};
 static constexpr size_t kNumTypes =
     sizeof(kTypeStrings) / sizeof(kTypeStrings[0]);
 static constexpr size_t kTypeBits[] = {
@@ -317,6 +320,7 @@ static constexpr size_t kTypeBits[] = {
     8 * sizeof(int8_t),
     4 /* Q4_0Stream, actually 4.5 */,
     4 /* MxFp4Stream, actually 4.25 */,
+    0 /* reserved */,
 };
 
 static inline bool EnumValid(Type type) {
@@ -376,17 +380,20 @@ constexpr bool IsCompressed() {
          hwy::IsSame<hwy::RemoveCvRef<Packed>, NuqStream>() ||
          hwy::IsSame<hwy::RemoveCvRef<Packed>, I8Stream>() ||
          hwy::IsSame<hwy::RemoveCvRef<Packed>, Q4_0Stream>() ||
-         hwy::IsSame<hwy::RemoveCvRef<Packed>, MxFp4Stream>();
+         hwy::IsSame<hwy::RemoveCvRef<Packed>, MxFp4Stream>()
+         ;
 }
 
 static inline bool IsCompressed(Type type) {
   return type == Type::kSFP || type == Type::kNUQ || type == Type::kI8 ||
-         type == Type::kQ4_0 || type == Type::kMXFP4;
+         type == Type::kQ4_0 || type == Type::kMXFP4
+         ;
 }
 
 static inline bool IsPacked(Type type) {
   return type == Type::kNUQ || type == Type::kI8 || type == Type::kQ4_0 ||
-         type == Type::kMXFP4;
+         type == Type::kMXFP4
+         ;
 }
 
 static inline bool SupportsPointerArithmetic(Type type) {

evals/model_health_test.cc

@@ -44,6 +44,7 @@
 #include "evals/cross_entropy.h"
 #include "gemma/configs.h"
 #include "gemma/gemma.h"
+#include "util/test_util.h"
 #include "hwy/base.h"
 #include "hwy/tests/hwy_gtest.h"
 
@@ -485,7 +486,7 @@ TEST_F(ModelHealthTest, DeterministicGeneration) {
 }  // namespace gcpp
 
 int main(int argc, char** argv) {
-  testing::InitGoogleTest(&argc, argv);
+  gcpp::InternalInitTest();
   gcpp::ModelHealthTest::InitEnv(argc, argv);
   int ret = RUN_ALL_TESTS();
   gcpp::ModelHealthTest::DeleteEnv();

gemma/activations.h

@@ -801,7 +801,8 @@ struct Activations {
                                    config.model_dim, allocator)),
           mla_o_in(MatFactory("mla_o_in",
                               mla_dims.o_in_dim > 0 ? batch_size : 0,
-                              mla_dims.o_in_dim, allocator)) {
+                              mla_dims.o_in_dim, allocator))
+          {
       moe_C1.AllocateAndAttachRowPtrs(row_ptrs);
       moe_C2.AllocateAndAttachRowPtrs(row_ptrs);
       ffw_expert_in.AllocateAndAttachRowPtrs(row_ptrs);

gemma/configs.cc

@@ -481,7 +481,7 @@ static LayerConfig LayerConfigGemma4_26B_MoE_LM(size_t model_dim) {
 static ModelConfig ConfigGemma4_26B_MoE() {
   ModelConfig config = ConfigBaseGemmaV4();
   config.display_name = "Gemma4_26B_MoE";
-  config.final_cap = 0.0f;
+  config.final_cap = 30.0f;
   config.att_cap = 0.0f;
   config.model = Model::GEMMA4_26B_MOE;
   config.wrapping = PromptWrapping::GEMMA_IT;

gemma/gemma4_moe.cc

@@ -48,6 +48,7 @@
 #include "gemma/attention.h"  // includes highway.h
 #include "gemma/tiled_attention.h"
 #include "gemma/gemma-inl.h"
+#include "ops/fast_ops-inl.h"
 #include "ops/ops-inl.h"
 
 HWY_BEFORE_NAMESPACE();
@@ -471,9 +472,12 @@ void Gemma4MoETransformerLayer(size_t num_tokens, size_t layer_idx,
                      /*is_attention=*/true, env.ctx);
 
   // Dual-Path FFW
-  pre_norm(
-      layer.pre_ffw2_ns.HasPtr() ? layer.pre_ffw2_ns : layer.pre_ffw_norm_scale,
-      activations.pre_ffw_rms_out);
+  const MatPtr& shared_norm = layer.pre_ffw2_ns.HasPtr() ? layer.pre_ffw2_ns : layer.pre_ffw_norm_scale;
+  const MatPtr& moe_norm = layer.pre_ffw_norm_scale;
+  const MatPtr& shared_post_norm = layer.post_ffw2_ns;
+  const MatPtr& moe_post_norm = layer.post_ffw1_ns;
+
+  pre_norm(shared_norm, activations.pre_ffw_rms_out);
 
   // Shared MLP Path
   FFWNoVit(layer, activations, env);  // writes to activations.ffw_out
@@ -487,17 +491,17 @@ void Gemma4MoETransformerLayer(size_t num_tokens, size_t layer_idx,
     }
   }
 
-  if (layer.post_ffw2_ns.HasPtr()) {
-    rms_norm_inplace(layer.post_ffw2_ns, activations.attention.att_sums);
+  if (shared_post_norm.HasPtr()) {
+    rms_norm_inplace(shared_post_norm, activations.attention.att_sums);
   }
 
   // MoE Path
-  pre_norm(layer.pre_ffw_norm_scale, activations.pre_ffw_rms_out);
+  pre_norm(moe_norm, activations.pre_ffw_rms_out);
 
   Gemma4MoE::MoEFFW(layer, activations, env);  // writes to activations.ffw_out
 
-  if (layer.post_ffw1_ns.HasPtr()) {
-    rms_norm_inplace(layer.post_ffw1_ns, activations.ffw_out);
+  if (moe_post_norm.HasPtr()) {
+    rms_norm_inplace(moe_post_norm, activations.ffw_out);
   }
 
   // Combine & Final Norm (Fix for dual-path combination)

gemma/tensor_info.h

@@ -68,6 +68,9 @@ struct TensorInfo {
   // If false, then [10, 20, 30] -> [10*20, 30] and [30] -> [1, 30].
   // If true, then [10, 20, 30] -> [10, 20*30] and [30] -> [1, 30].
   bool cols_take_extra_dims = false;
+  // Optional pre-computed scale (e.g. for kW2_UL weights from QAFT).
+  // If > 0.0, used directly instead of re-estimating via ScaleWeightsW2UL.
+  float scale = 0.0f;
 };
 
 // Collapses/expands the tensor dims into 2D extents, which may be 0, 0 for

gemma/weights.cc

@@ -20,6 +20,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#include <cstring>
 #include <mutex>  // NOLINT
 #include <string>
 #include <vector>
@@ -129,7 +130,10 @@ static void SplitPackedMatrix(MatPtr& parent, size_t split_row, MatPtr& w1,
 void LayerWeightsPtrs::SplitW1() {
   // Used for Gemma layers; FFWVit uses different tensors.
   if (layer_config.type == LayerAttentionType::kVit) return;
-  if (layer_config.IsMoE()) return;
+  if (layer_config.IsMoE() && !gating_einsum_w.HasPtr() &&
+      !gating_einsum_w1.HasPtr()) {
+    return;
+  }
 
   // Files have both or neither of w1 and w2.
   HWY_ASSERT(gating_einsum_w1.HasPtr() == gating_einsum_w2.HasPtr());
@@ -543,6 +547,18 @@ void LayerWeightsPtrs::Fixup(Model model, std::vector<MatOwner>& mat_owners,
       const size_t elem_bytes = qkv_einsum_w2.ElementBytes();
       const size_t old_row_bytes = old_stride * elem_bytes;
       const size_t kv_heads = layer_config.kv_heads;
+      const size_t qkv_dim = layer_config.qkv_dim;
+
+      // In Gemma 4 global layers, attention_k_eq_v is true (K0 == V0).
+      // If already interleaved by exporter: [K0, V0, K1, V1], slice 0 == slice 1.
+      // If not interleaved: [K0, K1, V0, V1], slice 0 (K0) != slice 1 (K1).
+      const uint8_t* slice0 = qkv_einsum_w2.RowBytes(0);
+      const uint8_t* slice1 = qkv_einsum_w2.RowBytes(qkv_dim);
+      if (std::memcmp(slice0, slice1, old_row_bytes) == 0) {
+        // Exporter already emitted interleaved layout; skip fixup.
+        return;
+      }
+
       const size_t total_bytes = qkv_einsum_w2.Rows() * old_row_bytes;
       hwy::AlignedFreeUniquePtr<uint8_t[]> tmp =
           hwy::AllocateAligned<uint8_t>(total_bytes);
@@ -556,7 +572,6 @@ void LayerWeightsPtrs::Fixup(Model model, std::vector<MatOwner>& mat_owners,
       }
 
       const size_t new_row_bytes = qkv_einsum_w2.Cols() * elem_bytes;
-      const size_t qkv_dim = layer_config.qkv_dim;
       const uint8_t* src_ptr = tmp.get();
       for (size_t i = 0; i < kv_heads; ++i) {
         for (size_t row = 0; row < qkv_dim; ++row) {
@@ -676,12 +691,12 @@ WeightsPtrs::Mode weights_internal::ChooseMode(uint64_t file_bytes,
     map = Tristate::kFalse;
   }
 
-  // Disable mapping if not padded to the base page size.
-  if (file_bytes % allocator.BasePageBytes() != 0) {
+  // Kernels can read multiple vectors ahead, so retain blob alignment even
+  // though mapping itself does not require a page-aligned file length.
+  if (file_bytes % kBlobAlign != 0) {
     if (map == Tristate::kTrue) {  // Only complain if explicitly requested.
-      HWY_WARN("Unable to map non-padded file (%zu, %zu), reading instead.",
-               static_cast<size_t>(file_bytes >> 10),
-               allocator.BasePageBytes());
+      HWY_WARN("File size %zu is not a multiple of %zu bytes, reading instead.",
+               static_cast<size_t>(file_bytes), kBlobAlign);
     }
     map = Tristate::kFalse;
   }

gemma/weights_test.cc

@@ -219,5 +219,71 @@ TEST(WeightsTest, ExplicitSFPDisablesOnlyAutomaticMapping) {
             WeightsPtrs::Mode::kReadBF16);
 }
 
+TEST(WeightsTest, MapsBlobAlignedNonPageAlignedFile) {
+  ThreadingContext ctx = MakeContext();
+  const size_t file_bytes = ctx.allocator.BasePageBytes() + kBlobAlign;
+  ASSERT_EQ(file_bytes % kBlobAlign, 0);
+  ASSERT_NE(file_bytes % ctx.allocator.BasePageBytes(), 0);
+  std::vector<uint8_t> contents(file_bytes);
+  contents.front() = 1;
+  contents.back() = 2;
+
+  TemporaryBlob blob;
+  auto file = OpenFileOrAbort(blob.path(), "w+");
+  ASSERT_TRUE(file->Write(contents.data(), contents.size(), 0));
+  MapPtr mapped = file->Map();
+  ASSERT_NE(mapped, nullptr);
+  EXPECT_EQ(mapped[0], contents.front());
+  EXPECT_EQ(mapped[file_bytes - 1], contents.back());
+
+  InferenceArgs inference;
+  LoaderArgs loader("", "");
+  loader.map = Tristate::kTrue;
+  loader.to_bf16 = Tristate::kFalse;
+  EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference,
+                                         ctx.allocator),
+            WeightsPtrs::Mode::kMap);
+}
+
+TEST(WeightsTest, MappingRequiresBlobAlignment) {
+  ThreadingContext ctx = MakeContext();
+  InferenceArgs inference;
+  LoaderArgs loader("", "");
+  loader.to_bf16 = Tristate::kFalse;
+
+  // Large enough for the automatic mapping heuristic to select kMap.
+  const uint64_t file_mib = ctx.allocator.TotalMiB() / 3 + 1;
+  const uint64_t page_bytes = ctx.allocator.BasePageBytes();
+  const uint64_t page_aligned_bytes =
+      hwy::RoundUpTo(file_mib << 20, page_bytes);
+
+  for (const Tristate map : {Tristate::kDefault, Tristate::kTrue}) {
+    loader.map = map;
+    for (const uint64_t file_bytes :
+         {page_aligned_bytes, page_aligned_bytes + kBlobAlign}) {
+      EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference,
+                                             ctx.allocator),
+                WeightsPtrs::Mode::kMap)
+          << "file_bytes=" << file_bytes;
+    }
+    // Cover both an almost-full page (Jan's 4090-byte example) and a file
+    // whose final page has plenty of space but lacks blob alignment.
+    for (const uint64_t file_bytes :
+         {page_bytes - 6, page_bytes + 1, page_aligned_bytes - 6,
+          page_aligned_bytes + 1}) {
+      EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference,
+                                             ctx.allocator),
+                WeightsPtrs::Mode::kRead)
+          << "file_bytes=" << file_bytes;
+    }
+  }
+
+  loader.map = Tristate::kTrue;
+  loader.to_bf16 = Tristate::kTrue;
+  EXPECT_EQ(weights_internal::ChooseMode(page_aligned_bytes + 1, loader,
+                                         inference, ctx.allocator),
+            WeightsPtrs::Mode::kReadBF16);
+}
+
 }  // namespace
 }  // namespace gcpp

io/blob_store.cc

@@ -35,15 +35,9 @@ namespace gcpp {
 
 static_assert(HWY_IS_LITTLE_ENDIAN, "Assumes little endian");
 
-// Each blob offset is a multiple of this, an upper bound on SVE vectors and
-// usually also larger than L2 cache lines. This is useful when memory mapping
-// the entire file, because offset alignment then determines the alignment of
-// the blob in memory. Aligning each blob to the (largest) page size would be
-// too wasteful, see `kEndAlign`.
-constexpr size_t kBlobAlign = 256;  // test also hard-codes this value
-
-// Linux mmap requires the file to be a multiple of the (base) page size, which
-// can be up to 64 KiB on Arm. Apple uses 16 KiB, most others use 4 KiB.
+// Pad newly written files to cover base page sizes up to 64 KiB on Arm.
+// Apple uses 16 KiB, most others use 4 KiB. Mapping also accepts older files
+// padded only to kBlobAlign.
 constexpr size_t kEndAlign = 64 * 1024;
 
 constexpr size_t kU128Bytes = sizeof(hwy::uint128_t);

io/blob_store.h

@@ -34,6 +34,11 @@
 
 namespace gcpp {
 
+// Blob offsets and padded sizes are multiples of this, an upper bound on SVE
+// vector bytes and usually also larger than L2 cache lines. This preserves blob
+// alignment when mapping the entire file without requiring page-sized padding.
+constexpr size_t kBlobAlign = 256;  // test also hard-codes this value
+
 // One blob's extents within the file.
 struct BlobRange {
   uint64_t End() const { return offset + bytes; }