From 1e2fe4dad8c33bdd703078a5bd4150c2e91d7d30 Mon Sep 17 00:00:00 2001 From: Mikyx-1 Date: Sun, 23 Aug 2026 15:12:00 +0700 Subject: [PATCH 1/4] Compress the embedding table to SFP at load time `c_embedding` is capped at 16 bits by `.min_size = Type::kBF16` in tensor_info.cc, so it stays BF16 even in an SFP model where every other large tensor is 8-bit. For 2.0-2b-pt-sfp.sbs that single 256000x2304 tensor is 1.18 GB, 37% of the file. Add `--sfp_embedding` (default off) to compress it to SFP while loading, as suggested by google/gemma.cpp#164. This halves its footprint and, for models with tied input/output embeddings, the weight bandwidth of the per-token logits MatMul. Models with a separate output head receive only the memory benefit. The read paths already handle SFP: `EmbedToken` seeks by `Stride()`, and `CallMatMul` dispatches to matmul_static_sfp. `MakeBatches` reads directly into destination rows, so flagged tensors bypass it and are read via `ReadAllToSFP`. That reads 4 MiB row chunks rather than staging the whole tensor, which would raise peak RSS by more than the conversion saves. The file is read twice because SFP encodes a limited range of magnitudes, hence the per-tensor scale must be known before encoding; the second read comes from the OS cache. Conversion requires owned memory, so an explicit SFP request disables automatic mapping. An explicit `--map=1` still takes precedence and warns that SFP embedding conversion is ignored. Validate the source blob size before chunked reads so inconsistent metadata cannot cross into an adjacent blob. Add focused tests for F32 and BF16 conversion, scale propagation, final partial chunks, malformed blob sizes, and mapping-mode selection. Measured on an M2 (4 workers), 2.0-2b-pt-sfp.sbs, Mode::kRead: Gen.EmbeddingMatmul 654 -> 468 us/token (-28.5%) decode 11.21 -> 12.42 tok/s (+10.8%, median of 3) max RSS 3482 -> 2928 MiB (-554 MiB) cross entropy/byte 1.372889 -> 1.368894 (no degradation) Startup does not regress: ReadBatches drops more than ReadAllToSFP adds, because MakeBatches issues row-wise I/O for this padded 256000-row tensor. --- BUILD.bazel | 22 +++- CMakeLists.txt | 2 + gemma/gemma_args.h | 8 ++ gemma/weights.cc | 244 +++++++++++++++++++++++++++++++++++---- gemma/weights_internal.h | 61 ++++++++++ gemma/weights_test.cc | 231 ++++++++++++++++++++++++++++++++++++ util/zones.cc | 4 + util/zones.h | 2 + 8 files changed, 548 insertions(+), 26 deletions(-) create mode 100644 gemma/weights_internal.h create mode 100644 gemma/weights_test.cc diff --git a/BUILD.bazel b/BUILD.bazel index 77c5090a..ed616a39 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -349,7 +349,10 @@ cc_library( cc_library( name = "weights", srcs = ["gemma/weights.cc"], - hdrs = ["gemma/weights.h"], + hdrs = [ + "gemma/weights.h", + "gemma/weights_internal.h", + ], deps = [ ":configs", ":gemma_args", @@ -379,6 +382,23 @@ cc_test( ], ) +cc_test( + name = "weights_test", + srcs = ["gemma/weights_test.cc"], + deps = [ + ":gemma_args", + ":mat", + ":threading_context", + ":weights", + "//compression:compress", + "//compression:types", + "//io", + "//io:blob_store", + "@googletest//:gtest_main", # buildcleaner: keep + "@highway//:hwy", + ], +) + # For building all tests in one command, so we can test several. test_suite( name = "ops_tests", diff --git a/CMakeLists.txt b/CMakeLists.txt index 059145fe..f10c6be8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -206,6 +206,7 @@ set(SOURCES gemma/vit.h gemma/weights.cc gemma/weights.h + gemma/weights_internal.h io/blob_store.cc io/blob_store.h io/fields.cc @@ -383,6 +384,7 @@ set(GEMMA_TEST_FILES deepseek/deepseek_test.cc gemma/gemma_args_test.cc gemma/tensor_info_test.cc + gemma/weights_test.cc io/blob_store_test.cc io/fields_test.cc ops/bench_matmul.cc diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index 8ba8f0d7..08b68490 100644 --- a/gemma/gemma_args.h +++ b/gemma/gemma_args.h @@ -53,6 +53,7 @@ struct LoaderArgs : public ArgsBase { Path weights; // weights file location Tristate map; Tristate to_bf16; + Tristate sfp_embedding; Tristate wrapping; template @@ -65,6 +66,13 @@ struct LoaderArgs : public ArgsBase { "Enable memory-mapping? -1 = auto, 0 = no, 1 = yes."); visitor(to_bf16, "to_bf16", Tristate::kDefault, "Convert weights to bf16? -1 = auto, 0 = no, 1 = yes."); + visitor(sfp_embedding, "sfp_embedding", Tristate::kDefault, + "Compress the embedding table to SFP while loading? This halves\n" + " its footprint and, for models with tied input/output\n" + " embeddings, the bandwidth of the per-token logits MatMul,\n" + " in exchange for 8-bit precision on that tensor. Enabling this\n" + " disables automatic mapping; explicit --map=1 still wins.\n" + " -1 = auto (currently off), 0 = no, 1 = yes."); visitor(wrapping, "wrapping", Tristate::kDefault, "Enable prompt wrapping? Specify 0 for pre-2025 format PT models."); } diff --git a/gemma/weights.cc b/gemma/weights.cc index 535f8a88..4bd1186f 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -29,19 +29,22 @@ #include "gemma/configs.h" #include "gemma/gemma_args.h" #include "gemma/model_store.h" +#include "gemma/weights_internal.h" +#include "hwy/base.h" +#include "hwy/highway.h" +#include "hwy/profiler.h" #include "io/blob_store.h" #include "util/mat.h" #include "util/threading_context.h" #include "util/zones.h" -#include "hwy/base.h" -#include "hwy/highway.h" -#include "hwy/profiler.h" // TODO: move into foreach_target #include "compression/compress-inl.h" namespace gcpp { +using weights_internal::TensorToRead; + static std::mutex g_mat_owners_mutex; // Copies att_weights from `attn_vec_einsum_w`. @@ -657,13 +660,20 @@ std::vector WeightsPtrs::AddTensorDataToWriter( } // Decides whether to read or map based on heuristics and user override. -static WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, - const LoaderArgs& loader, - const InferenceArgs& inference, - const Allocator& allocator) { +WeightsPtrs::Mode weights_internal::ChooseMode(uint64_t file_bytes, + const LoaderArgs& loader, + const InferenceArgs& inference, + const Allocator& allocator) { Tristate to_bf16 = loader.to_bf16; Tristate map = loader.map; + // An explicit request to convert the embedding requires owned memory. Do + // not let the automatic mapping heuristic override that request. An + // explicit --map=1 is still honored and diagnosed in ReadFromBlobs. + if (loader.sfp_embedding == Tristate::kTrue && map == Tristate::kDefault) { + map = Tristate::kFalse; + } + // Disable mapping if not padded to the base page size. if (file_bytes % allocator.BasePageBytes() != 0) { if (map == Tristate::kTrue) { // Only complain if explicitly requested. @@ -706,18 +716,6 @@ static WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, : WeightsPtrs::Mode::kRead; } -struct TensorToRead { - MatPtr* mat; - BlobRange range; - // Some tensors opt out of padding via kPacked flags. - MatPadding padding; - - // only for kReadBF16 - bool keep_type = false; - Type prev_type; - size_t prev_packed_bytes = 0; -}; - // Allocates multiple in parallel and binds to NUMA nodes. static void AllocateAndBindAll(std::vector& tensors, const WeightsPtrs::Mode mode, @@ -735,8 +733,16 @@ static void AllocateAndBindAll(std::vector& tensors, tensor.prev_type = mat.GetType(); tensor.prev_packed_bytes = mat.PackedBytes(); - // We only care about MatMul inputs; skip F32 or small tensors. - if (tensor.prev_type == Type::kF32 || mat.Rows() < 1024) { + // Only worthwhile from 16/32-bit types; the others are already <= 8 + // bits, and NUQ is smaller than SFP. + if (tensor.to_sfp && tensor.prev_type != Type::kF32 && + tensor.prev_type != Type::kBF16) { + tensor.to_sfp = false; + } + if (tensor.to_sfp) { + mat.SetType(Type::kSFP); + // We only care about MatMul inputs; skip F32 or small tensors. + } else if (tensor.prev_type == Type::kF32 || mat.Rows() < 1024) { tensor.keep_type = true; tensor.padding = MatPadding::kPacked; // single I/O for simplicity } else if (mode == WeightsPtrs::Mode::kReadBF16) { @@ -836,6 +842,172 @@ static void ReadAllToBF16(const std::vector& tensors, }); } +// Tensors flagged `to_sfp`, in any mode that reads rather than maps: + +// Number of rows to read per parallel task. Rows are grouped so that the reads +// are large enough to be efficient, but small enough that the transient buffers +// are negligible: staging the whole tensor would defeat the purpose of +// compressing it. +static size_t RowsPerChunk(size_t row_bytes) { + constexpr size_t kTargetBytes = 4 * 1024 * 1024; + return HWY_MAX(size_t{1}, kTargetBytes / HWY_MAX(size_t{1}, row_bytes)); +} + +// Holds the per-task buffers for one chunk of `rows_per_chunk` rows, so that +// both passes below can reuse the same code. +template +class Chunk { + public: + Chunk(size_t cols, size_t rows_per_chunk) + : cols_(cols), rows_per_chunk_(rows_per_chunk) { + const hwy::HWY_NAMESPACE::ScalableTag df; + const size_t NF = hwy::HWY_NAMESPACE::Lanes(df); + buf_ = hwy::AllocateAligned(rows_per_chunk * cols * sizeof(T)); + // `DecompressAndZeroPad` writes whole vectors, and `compress-inl.h` + // requires up to two of them beyond the requested count. + raw_ = hwy::AllocateAligned( + hwy::RoundUpTo(rows_per_chunk * cols, NF) + 4 * NF); + HWY_ASSERT(buf_ && raw_); + } + + // Reads rows `[begin, end)` from the file, where the tensor is stored as + // type `T` and packed, hence `mat.Stride()` does not apply. Returns the + // decompressed values, ignoring `mat.Scale()`. + float* Decompress(const TensorToRead& tensor, const BlobReader& reader, + size_t begin, size_t end) { + HWY_DASSERT(end - begin <= rows_per_chunk_); + const size_t row_bytes = cols_ * sizeof(T); + const size_t num = (end - begin) * cols_; + HWY_ASSERT(reader.file().Read(tensor.range.offset + begin * row_bytes, + num * sizeof(T), buf_.get())); + + // Rows are contiguous in both source and destination, hence decompress the + // entire chunk at once: this is faster, and prevents the zero padding from + // overwriting the start of the next row. + const hwy::HWY_NAMESPACE::ScalableTag df; + const PackedSpan packed{HWY_RCAST_ALIGNED(T*, buf_.get()), num}; + HWY_NAMESPACE::DecompressAndZeroPad(df, packed, 0, raw_.get(), num); + return raw_.get(); + } + + private: + size_t cols_; + size_t rows_per_chunk_; + hwy::AlignedFreeUniquePtr buf_; + hwy::AlignedFreeUniquePtr raw_; +}; + +// Returns the largest magnitude in the tensor, ignoring `mat.Scale()`. +template +static float MaxAbs(const TensorToRead& tensor, const BlobReader& reader, + size_t rows_per_chunk, size_t num_chunks, + ThreadingContext& ctx) { + const MatPtr& mat = *tensor.mat; + const size_t rows = mat.Rows(); + const size_t cols = mat.Cols(); + // Indexed by chunk rather than by worker, so that we need not know how many + // workers `ParallelFor` will use. + std::vector chunk_max(num_chunks, 0.0f); + + ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + const float* raw = + buffers.Decompress(tensor, reader, begin, end); + + float maxabs = 0.0f; + for (size_t i = 0; i < (end - begin) * cols; ++i) { + maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); + } + chunk_max[chunk] = maxabs; + }); + + float maxabs = 0.0f; + for (const float m : chunk_max) maxabs = HWY_MAX(maxabs, m); + return maxabs; +} + +// Reads the tensor as stored in the file (type `T`) and compresses it into +// `mat`, whose type `AllocateAndBindAll` already changed to `Type::kSFP`. +// The file is read twice because `SfpStream` encodes a limited range of +// magnitudes, hence we may need a per-tensor scale, which requires knowing the +// largest magnitude before encoding anything. The second read is typically +// served from the OS cache. +template +static void CompressToSFP(const TensorToRead& tensor, const BlobReader& reader, + ThreadingContext& ctx) { + MatPtr& mat = *tensor.mat; + const size_t rows = mat.Rows(); + const size_t cols = mat.Cols(); + const size_t rows_per_chunk = RowsPerChunk(cols * sizeof(T)); + const size_t num_chunks = hwy::DivCeil(rows, rows_per_chunk); + + const float prev_scale = mat.Scale(); + const float maxabs = + MaxAbs(tensor, reader, rows_per_chunk, num_chunks, ctx) * prev_scale; + const float scale = + (maxabs <= SfpStream::kMax) ? 1.0f : maxabs / SfpStream::kMax; + const float mul = prev_scale / scale; + + ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + float* raw = buffers.Decompress(tensor, reader, begin, end); + + if (mul != 1.0f) { + for (size_t i = 0; i < (end - begin) * cols; ++i) { + // Clamp because rounding may still exceed `kMax`. + const float magn = + HWY_MIN(SfpStream::kMax, hwy::ScalarAbs(raw[i] * mul)); + raw[i] = hwy::ScalarCopySign(magn, raw[i]); + } + } + + // Row by row because destination rows are padded, whereas `raw` + // is not. This is safe because `SfpStream` is a per-value + // encoding. + CompressPerThread tls; // unused by SFP, which is stateless + for (size_t r = begin; r < end; ++r) { + const PackedSpan row{ + HWY_RCAST_ALIGNED(SfpStream*, mat.RowBytes(r)), cols}; + HWY_NAMESPACE::Compress(raw + (r - begin) * cols, cols, tls, + row, + /*packed_ofs=*/0); + } + }); + + mat.SetScale(scale); +} + +void weights_internal::ReadAllToSFP(const std::vector& tensors, + const BlobReader& reader, + ThreadingContext& ctx) { + PROFILER_ZONE("Startup.Weights.ReadAllToSFP"); + // Usually a single (large) tensor, hence parallelize within, not across. + for (const TensorToRead& tensor : tensors) { + // CompressToSFP derives read sizes from the tensor shape. Ensure those + // reads cannot cross the blob boundary if metadata is inconsistent. + HWY_ASSERT_M(tensor.range.bytes == tensor.prev_packed_bytes, + tensor.mat->Name()); + switch (tensor.prev_type) { + case Type::kF32: + CompressToSFP(tensor, reader, ctx); + break; + case Type::kBF16: + CompressToSFP(tensor, reader, ctx); + break; + default: + HWY_ABORT("Unsupported type %s", TypeName(tensor.prev_type)); + } + } +} + // Mode == kRead: static std::vector MakeBatches( @@ -934,13 +1106,20 @@ static MapPtr MapOrReadAll(std::vector& tensors, AllocateAndBindAll(tensors, *mode, mat_owners, ctx); } + // `MakeBatches` and `ReadAllToBF16` read into the destination rows, hence + // tensors that require a compression pass are handled separately. + std::vector to_sfp, rest; + for (const TensorToRead& tensor : tensors) { + (tensor.to_sfp ? to_sfp : rest).push_back(tensor); + } + if (!to_sfp.empty()) weights_internal::ReadAllToSFP(to_sfp, reader, ctx); + if (*mode == WeightsPtrs::Mode::kReadBF16) { - ReadAllToBF16(tensors, reader, ctx); + ReadAllToBF16(rest, reader, ctx); return MapPtr(); } - const std::vector batches = - MakeBatches(tensors, reader.file_bytes()); + const std::vector batches = MakeBatches(rest, reader.file_bytes()); ReadBatches(reader, batches, ctx); return MapPtr(); } @@ -973,7 +1152,22 @@ WeightsPtrs::Mode WeightsPtrs::ReadFromBlobs(const ModelStore& model, HWY_ABORT("Tensor %s is required but not found in file.", t.mat.Name()); }); - Mode mode = ChooseMode(reader.file_bytes(), loader, inference, ctx.allocator); + Mode mode = weights_internal::ChooseMode(reader.file_bytes(), loader, + inference, ctx.allocator); + + // Compressing the input embedding to SFP halves its footprint. For models + // with tied input/output embeddings, it also halves the weight bandwidth of + // the per-token logits MatMul. + if (loader.sfp_embedding == Tristate::kTrue) { + if (mode == Mode::kMap) { + HWY_WARN("Cannot have sfp_embedding && map, ignoring sfp_embedding."); + } else { + for (TensorToRead& tensor : tensors) { + if (tensor.mat == &embedder_input_embedding) tensor.to_sfp = true; + } + } + } + mapped_ = MapOrReadAll(tensors, reader, &mode, mat_owners, ctx); { diff --git a/gemma/weights_internal.h b/gemma/weights_internal.h new file mode 100644 index 00000000..1d84ba6f --- /dev/null +++ b/gemma/weights_internal.h @@ -0,0 +1,61 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ +#define THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ + +#include +#include + +#include + +#include "compression/types.h" +#include "gemma/gemma_args.h" +#include "gemma/weights.h" +#include "io/blob_store.h" +#include "util/allocator.h" +#include "util/mat.h" +#include "util/threading_context.h" + +namespace gcpp { +namespace weights_internal { + +// Describes one tensor whose file bytes will be mapped, read, or converted. +// Kept in this internal header so loader behavior can be tested without a full +// model-sized weights file. +struct TensorToRead { + MatPtr* mat; + BlobRange range; + MatPadding padding; + + // Only for kReadBF16. + bool keep_type = false; + // Convert to Type::kSFP while loading. + bool to_sfp = false; + Type prev_type; + size_t prev_packed_bytes = 0; +}; + +WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, const LoaderArgs& loader, + const InferenceArgs& inference, + const Allocator& allocator); + +void ReadAllToSFP(const std::vector& tensors, + const BlobReader& reader, ThreadingContext& ctx); + +} // namespace weights_internal +} // namespace gcpp + +#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ diff --git a/gemma/weights_test.cc b/gemma/weights_test.cc new file mode 100644 index 00000000..8692ac87 --- /dev/null +++ b/gemma/weights_test.cc @@ -0,0 +1,231 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include +#include +#include +#include + +#include "compression/compress-inl.h" +#include "compression/types.h" +#include "gemma/gemma_args.h" +#include "gemma/weights_internal.h" +#include "gtest/gtest.h" +#include "hwy/base.h" +#include "io/blob_store.h" +#include "io/io.h" +#include "util/basics.h" +#include "util/mat.h" +#include "util/threading_context.h" + +namespace gcpp { +namespace { + +class TemporaryBlob { + public: + TemporaryBlob() { + const int fd = mkstemp(path_); + HWY_ASSERT(fd >= 0); + HWY_ASSERT(close(fd) == 0); + } + + ~TemporaryBlob() { unlink(path_); } + + Path path() const { return Path(path_); } + + private: + char path_[sizeof("/tmp/weights_test.sbs-XXXXXX")] = + "/tmp/weights_test.sbs-XXXXXX"; +}; + +static ThreadingContext MakeContext() { + ThreadingArgs args; + args.max_threads = 2; + args.pin = Tristate::kFalse; + args.bind = Tristate::kFalse; + return ThreadingContext(args); +} + +template +static float SourceValue(const T value) { + return static_cast(value); +} + +template <> +float SourceValue(const BF16 value) { + return hwy::F32FromBF16(value); +} + +template +static T MakeSourceValue(const float value) { + return static_cast(value); +} + +template <> +BF16 MakeSourceValue(const float value) { + return hwy::BF16FromF32(value); +} + +template +static void TestSFPConversion(const Type source_type, const float prev_scale, + const std::array& pattern) { + ThreadingContext ctx = MakeContext(); + + // Just over 4 MiB, so the conversion has a final partial chunk. + constexpr size_t kCols = 1024; + constexpr size_t kRows = 4 * 1024 * 1024 / (kCols * sizeof(T)) + 1; + const size_t num = kRows * kCols; + std::vector source(num); + for (size_t i = 0; i < num; ++i) { + source[i] = MakeSourceValue(pattern[i % pattern.size()]); + } + + TemporaryBlob blob; + { + BlobWriter writer(blob.path(), ctx); + writer.Add("embedding", source.data(), source.size() * sizeof(T)); + writer.Finalize(); + } + BlobReader reader(blob.path()); + const BlobRange* range = reader.Find("embedding"); + ASSERT_NE(range, nullptr); + + MatPtr mat("embedding", Type::kSFP, Extents2D(kRows, kCols)); + mat.SetScale(prev_scale); + MatOwner owner; + owner.AllocateFor(mat, ctx.allocator, MatPadding::kOdd); + + const weights_internal::TensorToRead tensor{ + .mat = &mat, + .range = *range, + .padding = MatPadding::kOdd, + .to_sfp = true, + .prev_type = source_type, + .prev_packed_bytes = source.size() * sizeof(T), + }; + weights_internal::ReadAllToSFP({tensor}, reader, ctx); + + float source_maxabs = 0.0f; + for (const T value : source) { + source_maxabs = + std::max(source_maxabs, std::abs(SourceValue(value) * prev_scale)); + } + const float expected_scale = + source_maxabs <= SfpStream::kMax ? 1.0f : source_maxabs / SfpStream::kMax; + EXPECT_FLOAT_EQ(mat.Scale(), expected_scale); + + const MatPtrT sfp(mat); + const std::array samples = { + 0, + 1, + kCols - 1, + (kRows - 1) * kCols - 1, + (kRows - 1) * kCols, + (kRows - 1) * kCols + 1, + num - 2, + num - 1, + }; + for (const size_t i : samples) { + const size_t row = i / kCols; + const size_t col = i % kCols; + const float expected = SourceValue(source[i]) * prev_scale; + const float actual = HWY_NAMESPACE::CompressTraits::ToFloatSlow( + sfp.Row(row)[col]) * + mat.Scale(); + const float tolerance = std::max(1E-6f, std::abs(expected) * 0.13f); + EXPECT_NEAR(actual, expected, tolerance) << "index " << i; + } +} + +TEST(WeightsTest, ReadBF16EmbeddingToSFPWithScaleAndPartialChunk) { + TestSFPConversion(Type::kBF16, 1.25f, + {-2.5f, -1.25f, -0.25f, 0.0f, 0.25f, 1.0f, 2.5f}); +} + +TEST(WeightsTest, ReadF32EmbeddingToSFPWithoutScaleAndPartialChunk) { + TestSFPConversion(Type::kF32, 1.0f, + {-1.75f, -1.0f, -0.125f, 0.0f, 0.125f, 1.0f, 1.75f}); +} + +TEST(WeightsTest, RejectsMismatchedEmbeddingBlobSize) { + ThreadingContext ctx = MakeContext(); + TemporaryBlob blob; + const std::array source = { + hwy::BF16FromF32(-1.0f), hwy::BF16FromF32(0.0f), hwy::BF16FromF32(1.0f), + hwy::BF16FromF32(2.0f)}; + { + BlobWriter writer(blob.path(), ctx); + writer.Add("embedding", source.data(), sizeof(source)); + writer.Finalize(); + } + BlobReader reader(blob.path()); + const BlobRange* range = reader.Find("embedding"); + ASSERT_NE(range, nullptr); + + MatPtr mat("embedding", Type::kSFP, Extents2D(2, 2)); + const weights_internal::TensorToRead tensor{ + .mat = &mat, + .range = *range, + .padding = MatPadding::kOdd, + .to_sfp = true, + .prev_type = Type::kBF16, + .prev_packed_bytes = sizeof(source) + sizeof(BF16), + }; + + EXPECT_DEATH(weights_internal::ReadAllToSFP({tensor}, reader, ctx), + "tensor.range.bytes == tensor.prev_packed_bytes"); +} + +TEST(WeightsTest, ExplicitSFPDisablesOnlyAutomaticMapping) { + ThreadingContext ctx = MakeContext(); + InferenceArgs inference; + LoaderArgs loader("", ""); + loader.to_bf16 = Tristate::kFalse; + + const uint64_t file_mib = ctx.allocator.TotalMiB() / 3 + 1; + const uint64_t file_bytes = hwy::RoundUpTo( + file_mib << 20, static_cast(ctx.allocator.BasePageBytes())); + + // Establish that the normal automatic heuristic would map this file. + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap); + + loader.sfp_embedding = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kRead); + + // An explicit mapping request still wins and is warned about by the loader. + loader.map = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap); + + // SFP embedding conversion composes with explicit conversion of other + // tensors to BF16. + loader.map = Tristate::kDefault; + loader.to_bf16 = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kReadBF16); +} + +} // namespace +} // namespace gcpp diff --git a/util/zones.cc b/util/zones.cc index 2c63f61c..9cd3475f 100644 --- a/util/zones.cc +++ b/util/zones.cc @@ -116,6 +116,8 @@ const char* ZoneName(Zones zone) { return "Ops.Softmax"; case Zones::kStartupWeightsReadAllToBF16: return "Startup.Weights.ReadAllToBF16"; + case Zones::kStartupWeightsReadAllToSFP: + return "Startup.Weights.ReadAllToSFP"; case Zones::kStartupWeightsReadBatches: return "Startup.Weights.ReadBatches"; default: @@ -207,6 +209,8 @@ const char* CallerName(Callers caller) { return "Ops.RMSNormNoScaleInplaceBatched"; case Callers::kReadAllToBF16: return "ReadAllToBF16"; + case Callers::kReadAllToSFP: + return "ReadAllToSFP"; case Callers::kReadBatches: return "ReadBatches"; case Callers::kSampleAndStream: diff --git a/util/zones.h b/util/zones.h index 01adb005..533182a9 100644 --- a/util/zones.h +++ b/util/zones.h @@ -66,6 +66,7 @@ enum class Zones { // Keep sorted kOpsRopeAndMulBy, kOpsSoftmax, kStartupWeightsReadAllToBF16, + kStartupWeightsReadAllToSFP, kStartupWeightsReadBatches, kNumZones // must be last }; @@ -116,6 +117,7 @@ enum class Callers { // Keep sorted kOpsRMSNormInplaceBatched, kOpsRMSNormNoScaleInplaceBatched, kReadAllToBF16, + kReadAllToSFP, kReadBatches, kSampleAndStream, kTensorStats, From 4b49a8fb34505090439a170296eaba71ea28cdda Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Fri, 28 Aug 2026 10:43:43 +0700 Subject: [PATCH 2/4] Use Highway scalar conversions in weights test --- gemma/weights_test.cc | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/gemma/weights_test.cc b/gemma/weights_test.cc index 8692ac87..78e810f1 100644 --- a/gemma/weights_test.cc +++ b/gemma/weights_test.cc @@ -64,22 +64,12 @@ static ThreadingContext MakeContext() { template static float SourceValue(const T value) { - return static_cast(value); -} - -template <> -float SourceValue(const BF16 value) { - return hwy::F32FromBF16(value); + return hwy::ConvertScalarTo(value); } template static T MakeSourceValue(const float value) { - return static_cast(value); -} - -template <> -BF16 MakeSourceValue(const float value) { - return hwy::BF16FromF32(value); + return hwy::ConvertScalarTo(value); } template From a075d4ecce68a67e5058fd221f4d5f9394166bbb Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Fri, 28 Aug 2026 11:07:57 +0700 Subject: [PATCH 3/4] Vectorize the embedding max-absolute scan --- gemma/weights.cc | 53 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/gemma/weights.cc b/gemma/weights.cc index 4bd1186f..31f72c13 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -909,21 +909,44 @@ static float MaxAbs(const TensorToRead& tensor, const BlobReader& reader, // workers `ParallelFor` will use. std::vector chunk_max(num_chunks, 0.0f); - ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, - Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { - GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); - const size_t begin = chunk * rows_per_chunk; - const size_t end = HWY_MIN(begin + rows_per_chunk, rows); - Chunk buffers(cols, rows_per_chunk); - const float* raw = - buffers.Decompress(tensor, reader, begin, end); - - float maxabs = 0.0f; - for (size_t i = 0; i < (end - begin) * cols; ++i) { - maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); - } - chunk_max[chunk] = maxabs; - }); + ParallelFor( + Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + const float* raw = buffers.Decompress(tensor, reader, begin, end); + + using DF = hwy::HWY_NAMESPACE::ScalableTag; + using VF = hwy::HWY_NAMESPACE::Vec; + const DF df; + const size_t NF = hwy::HWY_NAMESPACE::Lanes(df); + const size_t num = (end - begin) * cols; + VF max0 = hwy::HWY_NAMESPACE::Zero(df); + VF max1 = hwy::HWY_NAMESPACE::Zero(df); + size_t i = 0; + for (; i + 2 * NF <= num; i += 2 * NF) { + max0 = hwy::HWY_NAMESPACE::Max( + max0, + hwy::HWY_NAMESPACE::Abs(hwy::HWY_NAMESPACE::LoadU(df, raw + i))); + max1 = hwy::HWY_NAMESPACE::Max( + max1, hwy::HWY_NAMESPACE::Abs( + hwy::HWY_NAMESPACE::LoadU(df, raw + i + NF))); + } + if (i + NF <= num) { + max0 = hwy::HWY_NAMESPACE::Max( + max0, + hwy::HWY_NAMESPACE::Abs(hwy::HWY_NAMESPACE::LoadU(df, raw + i))); + i += NF; + } + float maxabs = hwy::HWY_NAMESPACE::ReduceMax( + df, hwy::HWY_NAMESPACE::Max(max0, max1)); + for (; i < num; ++i) { + maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); + } + chunk_max[chunk] = maxabs; + }); float maxabs = 0.0f; for (const float m : chunk_max) maxabs = HWY_MAX(maxabs, m); From f3407aff593e0f6fed6c7c0df41b870941fbd003 Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Fri, 28 Aug 2026 21:55:55 +0700 Subject: [PATCH 4/4] Use Highway namespace alias in weights --- gemma/weights.cc | 87 +++++++++++++++++++++++------------------------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/gemma/weights.cc b/gemma/weights.cc index 31f72c13..c699ac6a 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -43,6 +43,8 @@ namespace gcpp { +namespace hn = hwy::HWY_NAMESPACE; + using weights_internal::TensorToRead; static std::mutex g_mat_owners_mutex; @@ -309,7 +311,7 @@ static void HWY_MAYBE_UNUSED InitAttWeightsI8( hwy::AlignedFreeUniquePtr att_weights_tmp = hwy::AllocateAligned(model_dim * heads * qkv_dim); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, attn_vec_einsum_w.Span(), 0, attn_vec_einsum_w_tmp.get(), model_dim * heads * qkv_dim); @@ -373,7 +375,7 @@ static void HWY_MAYBE_UNUSED SplitW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(total_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, gating_einsum_w.Span(), 0, w_tmp.get(), total_size); @@ -430,7 +432,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(w1_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, qkv_einsum_w.Span(), 0, w_tmp.get(), w1_size); @@ -469,7 +471,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(total_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, qkv_einsum_w.Span(), 0, w_tmp.get(), total_size); @@ -777,7 +779,7 @@ static void MapAll(const std::vector& tensors, template static void DecompressToBF16(MatPtr& mat, const hwy::AlignedFreeUniquePtr& buf) { - hwy::HWY_NAMESPACE::ScalableTag dbf; + hn::ScalableTag dbf; const size_t cols = mat.Cols(); const size_t num_packed = CompressedArrayElements(mat.Extents().Area()); @@ -860,8 +862,8 @@ class Chunk { public: Chunk(size_t cols, size_t rows_per_chunk) : cols_(cols), rows_per_chunk_(rows_per_chunk) { - const hwy::HWY_NAMESPACE::ScalableTag df; - const size_t NF = hwy::HWY_NAMESPACE::Lanes(df); + const hn::ScalableTag df; + const size_t NF = hn::Lanes(df); buf_ = hwy::AllocateAligned(rows_per_chunk * cols * sizeof(T)); // `DecompressAndZeroPad` writes whole vectors, and `compress-inl.h` // requires up to two of them beyond the requested count. @@ -884,7 +886,7 @@ class Chunk { // Rows are contiguous in both source and destination, hence decompress the // entire chunk at once: this is faster, and prevents the zero padding from // overwriting the start of the next row. - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; const PackedSpan packed{HWY_RCAST_ALIGNED(T*, buf_.get()), num}; HWY_NAMESPACE::DecompressAndZeroPad(df, packed, 0, raw_.get(), num); return raw_.get(); @@ -909,44 +911,37 @@ static float MaxAbs(const TensorToRead& tensor, const BlobReader& reader, // workers `ParallelFor` will use. std::vector chunk_max(num_chunks, 0.0f); - ParallelFor( - Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, - Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { - GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); - const size_t begin = chunk * rows_per_chunk; - const size_t end = HWY_MIN(begin + rows_per_chunk, rows); - Chunk buffers(cols, rows_per_chunk); - const float* raw = buffers.Decompress(tensor, reader, begin, end); - - using DF = hwy::HWY_NAMESPACE::ScalableTag; - using VF = hwy::HWY_NAMESPACE::Vec; - const DF df; - const size_t NF = hwy::HWY_NAMESPACE::Lanes(df); - const size_t num = (end - begin) * cols; - VF max0 = hwy::HWY_NAMESPACE::Zero(df); - VF max1 = hwy::HWY_NAMESPACE::Zero(df); - size_t i = 0; - for (; i + 2 * NF <= num; i += 2 * NF) { - max0 = hwy::HWY_NAMESPACE::Max( - max0, - hwy::HWY_NAMESPACE::Abs(hwy::HWY_NAMESPACE::LoadU(df, raw + i))); - max1 = hwy::HWY_NAMESPACE::Max( - max1, hwy::HWY_NAMESPACE::Abs( - hwy::HWY_NAMESPACE::LoadU(df, raw + i + NF))); - } - if (i + NF <= num) { - max0 = hwy::HWY_NAMESPACE::Max( - max0, - hwy::HWY_NAMESPACE::Abs(hwy::HWY_NAMESPACE::LoadU(df, raw + i))); - i += NF; - } - float maxabs = hwy::HWY_NAMESPACE::ReduceMax( - df, hwy::HWY_NAMESPACE::Max(max0, max1)); - for (; i < num; ++i) { - maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); - } - chunk_max[chunk] = maxabs; - }); + ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + const float* raw = + buffers.Decompress(tensor, reader, begin, end); + + using DF = hn::ScalableTag; + using VF = hn::Vec; + const DF df; + const size_t NF = hn::Lanes(df); + const size_t num = (end - begin) * cols; + VF max0 = hn::Zero(df); + VF max1 = hn::Zero(df); + size_t i = 0; + for (; i + 2 * NF <= num; i += 2 * NF) { + max0 = hn::Max(max0, hn::Abs(hn::LoadU(df, raw + i))); + max1 = hn::Max(max1, hn::Abs(hn::LoadU(df, raw + i + NF))); + } + if (i + NF <= num) { + max0 = hn::Max(max0, hn::Abs(hn::LoadU(df, raw + i))); + i += NF; + } + float maxabs = hn::ReduceMax(df, hn::Max(max0, max1)); + for (; i < num; ++i) { + maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); + } + chunk_max[chunk] = maxabs; + }); float maxabs = 0.0f; for (const float m : chunk_max) maxabs = HWY_MAX(maxabs, m);