diff --git a/.gitignore b/.gitignore index 1c13032e..150d9252 100644 --- a/.gitignore +++ b/.gitignore @@ -1,11 +1,12 @@ # Build directories .cache/ -bazel-*/ +bazel-* build-*/ build/ # Python cache python/*/__pycache__ +python/__pycache__ # Model files *.sbs @@ -22,4 +23,11 @@ python/*/__pycache__ # Local development .env -.env.local \ No newline at end of file +.env.local +.venv + +# OS Files +.DS_Store + +# Temporary +.temp/ diff --git a/BUILD.bazel b/BUILD.bazel index 7a506341..51ac2859 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -594,7 +594,7 @@ cc_test( ":gemma_args", ":kv_cache", ":threading_context", - "//testing/base/public:gunit_main", + "@googletest//:gtest_main", # buildcleaner: keep "@highway//:hwy", ], ) diff --git a/DEVELOPERS.md b/DEVELOPERS.md index 5d70fdbb..43581cfe 100644 --- a/DEVELOPERS.md +++ b/DEVELOPERS.md @@ -96,8 +96,8 @@ https://github.com/keras-team/keras-nlp/blob/master/tools/gemma/export_gemma_to_ From Pytorch, use the following script to generate uncompressed weights: https://github.com/google/gemma.cpp/blob/dev/compression/convert_weights.py -For PaliGemma, use `python/convert_from_safetensors` to create an SBS file -directly. +For PaliGemma and T5Gemma S/S, use `python/convert_from_safetensors` to create +an SBS file directly. For other models, `gemma_export_main.py` is not yet open sourced. diff --git a/README.md b/README.md index 4c957f0e..d5532579 100644 --- a/README.md +++ b/README.md @@ -269,6 +269,41 @@ A tall tree stands in front of the building, and a window on the building is visible from the water. The water is green, and the sky is blue. ``` +### T5Gemma Encoder-Decoder Model + +This repository includes experimental support for the T5Gemma S/S +encoder-decoder model. Convert a local Hugging Face safetensors checkpoint to +SBS with: + +```sh +python3 python/convert_from_safetensors.py \ +--model_specifier=t5gemma-s-s \ +--load_path /path/to/t5gemma/model.safetensors \ +--tokenizer_file /path/to/t5gemma/tokenizer.model \ +--sbs_file t5gemma-s-s-it.sbs \ +--metadata_file t5gemma-s-s-it.csv +``` + +The default T5Gemma conversion writes an all-BF16 SBS file, which is useful for +Hugging Face parity checks and is currently the recommended path. To write a +smaller experimental mixed BF16/SFP file, add `--t5gemma_weight_type=sfp`. + +Then run: + +```sh +./gemma \ +--tokenizer /path/to/t5gemma/tokenizer.model \ +--weights t5gemma-s-s-it.sbs \ +--model t5gemma-s-s \ +--prompt "Hello" +``` + +The first supported runtime path is fresh seq2seq generation. Multi-turn reuse +of decoder KV cache with new encoder inputs is intentionally not supported yet. +Instruction-tuned T5Gemma checkpoints use the default instruction wrapping. For +base/pre-trained checkpoints or raw Hugging Face token parity checks, pass +`--wrapping=0` to use pre-trained wrapping instead. + ### Migrating to single-file format There is now a new format for the weights file, which is a single file that diff --git a/compression/python/BUILD.bazel b/compression/python/BUILD.bazel index 6cc65281..97a52e7c 100644 --- a/compression/python/BUILD.bazel +++ b/compression/python/BUILD.bazel @@ -1,5 +1,5 @@ load("@rules_cc//cc:cc_library.bzl", "cc_library") -load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test") +load("@rules_python//python:defs.bzl", "py_test") load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") package( diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc index 3568ad34..084a1167 100644 --- a/compression/python/compression_clif_aux.cc +++ b/compression/python/compression_clif_aux.cc @@ -51,6 +51,12 @@ HWY_BEFORE_NAMESPACE(); namespace gcpp { namespace HWY_NAMESPACE { +ThreadingArgs SingleThreadArgs() { + ThreadingArgs args; + args.max_lps = 1; + return args; +} + // Implementation for the currently compiled SIMD target. class SbsWriterImpl : public ISbsWriter { template @@ -91,7 +97,7 @@ class SbsWriterImpl : public ISbsWriter { public: SbsWriterImpl(const std::string& sbs_path) - : ctx_(ThreadingArgs()), writer_(gcpp::Path(sbs_path), ctx_) {} + : ctx_(SingleThreadArgs()), writer_(gcpp::Path(sbs_path), ctx_) {} void Insert(const char* name, F32Span weights, Type type, const TensorInfo& tensor_info) override { diff --git a/gemma/activations.h b/gemma/activations.h index 481926f2..6426ccbd 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -350,7 +350,8 @@ struct Activations { Activations(const RuntimeConfig& runtime_config, const ModelConfig& config, size_t batch_size, size_t seq_len, ThreadingContext& ctx, std::vector>& row_ptrs) - : layer_config(config.layer_configs[0]), + : layer_config(config.is_encoder_decoder ? config.decoder_layer_configs[0] + : config.layer_configs[0]), x(MatFactory("x", batch_size, config.model_dim, ctx.allocator)), x_bf(MatFactory("x_bf", batch_size, config.model_dim, ctx.allocator)), @@ -368,26 +369,56 @@ struct Activations { MatFactory("ffw_out", batch_size, config.model_dim, ctx.allocator)), max_workers(ctx.pools.MaxWorkers()), - s_ffw_in(config.num_layers, max_workers), - s_ffw_hidden(config.num_layers, max_workers), - s_ffw_out(config.num_layers, max_workers), + s_ffw_in(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_ffw_hidden(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_ffw_out(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), router_in(MatFactory("router_in", MoEBatchSize(layer_config, batch_size), config.model_dim, ctx.allocator)), router_logits( MatFactory("router_logits", MoEBatchSize(layer_config, batch_size), layer_config.NumExperts(), ctx.allocator)), - s_router_in(config.num_layers, max_workers), - s_router_logits(config.num_layers, max_workers), - s_expert_in(config.num_layers, max_workers), - s_expert_hidden(config.num_layers, max_workers), - s_expert_out(config.num_layers, max_workers), - s_w_expert_in1(config.num_layers, max_workers), - s_w_expert_in2(config.num_layers, max_workers), - s_w_expert_hidden(config.num_layers, max_workers), - s_w_gating_einsum_w1(config.num_layers, max_workers), - s_w_gating_einsum_w2(config.num_layers, max_workers), - s_w_linear_w(config.num_layers, max_workers), + s_router_in(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_router_logits(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_expert_in(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_expert_hidden(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_expert_out(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_expert_in1(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_expert_in2(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_expert_hidden(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_gating_einsum_w1(config.is_encoder_decoder + ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_gating_einsum_w2(config.is_encoder_decoder + ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_w_linear_w(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), attention_impl(runtime_config.attention_impl), attention_storage(config, layer_config, batch_size, seq_len, runtime_config, ctx.pools.MaxWorkers(), ctx.allocator, diff --git a/gemma/configs.cc b/gemma/configs.cc index 833c3142..66591323 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -495,6 +495,56 @@ static ModelConfig ConfigGemma4_26B_MoE() { return config; } +static ModelConfig ConfigBaseT5Gemma() { + ModelConfig config = ConfigNoSSM(); + config.att_cap = 50.0f; + config.final_cap = 30.0f; + config.eos_id = 1; + config.secondary_eos_id = 107; + return config; +} + +static LayerConfig LayerConfigT5GemmaS(size_t model_dim) { + LayerConfig config; + config.model_dim = model_dim; + config.ff_hidden_dim = 1024; + config.heads = 8; + config.kv_heads = 8; + config.qkv_dim = 64; + config.optimized_gating = false; + config.post_norm = PostNormType::Scale; + return config; +} + +static ModelConfig ConfigT5Gemma_S_S() { + ModelConfig config = ConfigBaseT5Gemma(); + config.display_name = "T5Gemma_S_S"; + config.model = Model::T5GEMMA_S_S; + config.wrapping = PromptWrapping::GEMMA_PT; + config.model_dim = 512; + config.vocab_size = kVocabSize; + config.max_seq_len = 8192; + LayerConfig layer_config = LayerConfigT5GemmaS(config.model_dim); + config.is_encoder_decoder = true; + config.encoder_num_layers = 8; + config.encoder_layer_configs = {config.encoder_num_layers, layer_config}; + config.encoder_attention_window_sizes = + RepeatedAttentionWindowSizes<8, 2>({4096, config.max_seq_len}); + config.decoder_num_layers = 8; + config.decoder_layer_configs = {config.decoder_num_layers, layer_config}; + config.decoder_attention_window_sizes = + RepeatedAttentionWindowSizes<8, 2>({4096, config.max_seq_len}); + + // TODO: Update users of `layer_configs` to route encoder-decoder models + // through the explicit encoder/decoder stacks above. + config.num_layers = 8; + config.layer_configs = {config.num_layers, layer_config}; + config.query_scale = QueryScaleType::SqrtKeySize; + config.attention_window_sizes = + RepeatedAttentionWindowSizes<8, 2>({4096, config.max_seq_len}); + return config; +} + static ModelConfig ConfigFromModel(Model model) { switch (model) { case Model::GEMMA2_2B: @@ -529,6 +579,8 @@ static ModelConfig ConfigFromModel(Model model) { return ConfigGemma3_27B_LM(); case Model::GEMMA4_26B_MOE: return ConfigGemma4_26B_MoE(); + case Model::T5GEMMA_S_S: + return ConfigT5Gemma_S_S(); default: HWY_ABORT("Model type %d unknown.", static_cast(model)); } @@ -570,6 +622,8 @@ const char* ModelPrefix(Model model) { return "gemma3-27b-lm"; case Model::GEMMA4_26B_MOE: return "gemma4-26b-moe"; + case Model::T5GEMMA_S_S: + return "t5gemma-s-s"; default: HWY_ABORT("Model type %d unknown.", static_cast(model)); } @@ -753,6 +807,14 @@ bool ModelConfig::OverwriteWithCanonical() { Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) { switch (layers) { + case 8: + if (layer_types & kDeducedT5Gemma) { + return Model::T5GEMMA_S_S; + } + HWY_WARN("Failed to deduce model type from %s, layer count %zu types %x.", + blob_path.path.c_str(), layers, layer_types); + return Model::UNKNOWN; + case 18: return Model::GEMMA3_270M; diff --git a/gemma/configs.h b/gemma/configs.h index a4d270d7..cda24a88 100644 --- a/gemma/configs.h +++ b/gemma/configs.h @@ -136,7 +136,7 @@ enum class PostQKType { static inline bool EnumValid(PostQKType type) { return static_cast(type) < - static_cast(PostNormType::kSentinel); + static_cast(PostQKType::kSentinel); } // FFW activation function. @@ -221,6 +221,8 @@ enum class Model { GEMMA3_12B_LM, GEMMA3_27B_LM, GEMMA4_26B_MOE, + // T5Gemma family - starting with S/S. + T5GEMMA_S_S, kSentinel, }; @@ -438,6 +440,14 @@ struct ModelConfig : public IFields { visitor(use_global_timescale); visitor(partial_rotary_factor); + visitor(is_encoder_decoder); + visitor(encoder_num_layers); + visitor(encoder_layer_configs); + visitor(encoder_attention_window_sizes); + visitor(decoder_num_layers); + visitor(decoder_layer_configs); + visitor(decoder_attention_window_sizes); + // Append new fields here, then update `python/configs.cc`. } @@ -465,6 +475,16 @@ struct ModelConfig : public IFields { attention_window_sizes[i] = new_max_seq_len; } } + for (size_t i = 0; i < encoder_attention_window_sizes.size(); ++i) { + if (encoder_attention_window_sizes[i] == max_seq_len) { + encoder_attention_window_sizes[i] = new_max_seq_len; + } + } + for (size_t i = 0; i < decoder_attention_window_sizes.size(); ++i) { + if (decoder_attention_window_sizes[i] == max_seq_len) { + decoder_attention_window_sizes[i] = new_max_seq_len; + } + } max_seq_len = new_max_seq_len; } @@ -494,10 +514,23 @@ struct ModelConfig : public IFields { for (const auto& layer_config : layer_configs) { num_heads = HWY_MAX(num_heads, layer_config.heads); } + for (const auto& layer_config : encoder_layer_configs) { + num_heads = HWY_MAX(num_heads, layer_config.heads); + } + for (const auto& layer_config : decoder_layer_configs) { + num_heads = HWY_MAX(num_heads, layer_config.heads); + } return num_heads; } size_t KVCacheCols() const { + if (is_encoder_decoder) { + size_t cols = 0; + for (const auto& lc : decoder_layer_configs) { + cols += lc.CacheLayerSize(); + } + return cols; + } size_t cols = 0; for (const auto& lc : layer_configs) { cols += lc.CacheLayerSize(); @@ -553,6 +586,17 @@ struct ModelConfig : public IFields { InternalModelConfig internal; bool use_global_timescale = false; // for Gemma 3 float partial_rotary_factor = 1.0f; // Fraction of dims with RoPE (0.25 for Gemma4 MoE). + + // Text encoder-decoder models such as T5Gemma have separate encoder and + // decoder stacks. Existing decoder-only models leave these fields at their + // defaults and continue to use `layer_configs`. + bool is_encoder_decoder = false; + uint32_t encoder_num_layers = 0; + std::vector encoder_layer_configs; + std::vector encoder_attention_window_sizes; + uint32_t decoder_num_layers = 0; + std::vector decoder_layer_configs; + std::vector decoder_attention_window_sizes; }; // Returns the sub-config for the ViT model of the PaliGemma model. @@ -562,6 +606,7 @@ enum DeducedLayerTypes { kDeducedViT = 2, kDeduced448 = 4, // For ViT, 448x448 resolution instead of 224x224. kDeducedKqNorm = 8, + kDeducedT5Gemma = 16, }; // layer_types is one or more of `DeducedLayerTypes`. diff --git a/gemma/configs_test.cc b/gemma/configs_test.cc index 5d649758..ec763e15 100644 --- a/gemma/configs_test.cc +++ b/gemma/configs_test.cc @@ -53,4 +53,27 @@ TEST(ConfigsTest, TestAttentionImpl) { ASSERT_EQ(GetAttentionImpl("invalid"), AttentionImpl::kFlash); } +TEST(ConfigsTest, T5GemmaKVCacheUsesDecoderLayers) { + ModelConfig config(Model::T5GEMMA_S_S, Type::kSFP, PromptWrapping::GEMMA_PT); + ASSERT_TRUE(config.is_encoder_decoder); + ASSERT_FALSE(config.decoder_layer_configs.empty()); + + const size_t expected_cols = config.decoder_layer_configs.size() * + config.decoder_layer_configs[0].CacheLayerSize(); + EXPECT_EQ(config.KVCacheCols(), expected_cols); +} + +TEST(ConfigsTest, DeduceT5GemmaSS) { + EXPECT_EQ(DeduceModel(Path("t5gemma-s-s.sbs"), 8, kDeducedT5Gemma), + Model::T5GEMMA_S_S); +} + +TEST(ConfigsTest, T5GemmaBF16Specifier) { + ModelConfig config("t5gemma-s-s-bf16-it"); + EXPECT_EQ(config.model, Model::T5GEMMA_S_S); + EXPECT_EQ(config.weight, Type::kBF16); + EXPECT_EQ(config.wrapping, PromptWrapping::GEMMA_IT); + EXPECT_TRUE(config.is_encoder_decoder); +} + } // namespace gcpp diff --git a/gemma/gemma.cc b/gemma/gemma.cc index d55dc3dd..1b6686ab 100644 --- a/gemma/gemma.cc +++ b/gemma/gemma.cc @@ -19,6 +19,7 @@ #include "gemma/gemma.h" #include +#include #include #include #include @@ -55,6 +56,7 @@ #include #include +#include #include #include "gemma/configs.h" @@ -145,6 +147,836 @@ static float EmbeddingScaling(size_t model_dim) { hwy::ConvertScalarTo(sqrtf(static_cast(model_dim)))); } +static HWY_INLINE void EmbedTokenFromWeights(int token, size_t x_row, + const ModelConfig& model_config, + const MatPtr& embedding, + MatStorageT& x) { + const size_t model_dim = model_config.model_dim; + const float emb_scaling = EmbeddingScaling(model_dim); + + HWY_DASSERT(token >= 0); + HWY_DASSERT(token < static_cast(model_config.vocab_size)); + + CallUpcasted(&embedding, [&](const auto* weights_t) { + // Using `Stride` to compute the offset works for both NUQ (because we use + // an offset and NUQ is never padded) and padded, because non-NUQ types are + // seekable, hence the offset can also skip any padding. + const size_t embedding_ofs = token * weights_t->Stride(); + HWY_ASSERT(weights_t->Cols() == model_dim); + const auto embedding_span = + MakeSpan(weights_t->Row(0), embedding_ofs + model_dim); + DecompressAndZeroPad(hn::ScalableTag(), embedding_span, + embedding_ofs, x.Row(x_row), model_dim); + MulByConst(emb_scaling * weights_t->Scale(), x.Row(x_row), model_dim); + }); +} + +static constexpr int kT5GemmaPadId = 0; +static constexpr int kT5GemmaBosId = 2; + +static void StreamAndUpdateEOS(size_t qi, size_t pos, int token, + float prob, const ModelConfig& config, + const RuntimeConfig& runtime_config, + QBatch& qbatch, bool update_pos, + hwy::BitSet4096<>& non_eos); + +static HWY_INLINE SampleFunc ChooseSampleFunc(const RuntimeConfig& runtime_config, + const AesCtrEngine& engine, + ThreadingContext& ctx); + +static void ValidateT5GemmaFreshGeneration(size_t pos, size_t prefix_end) { + if (pos != 0 || prefix_end != 0) { + HWY_ABORT( + "T5Gemma currently supports fresh seq2seq generation only; got pos=%zu " + "and prefix_end=%zu.", + pos, prefix_end); + } +} + +static void InitT5GemmaEncoderCache(const ModelConfig& config, + const PromptTokens& prompt, + T5GemmaEncoderCache& cache, + const Allocator& allocator) { + HWY_ASSERT(prompt.size() != 0); + cache.source_len = prompt.size(); + cache.hidden_states = + MatStorageT("t5_enc", Extents2D(cache.source_len, config.model_dim), + allocator, MatPadding::kOdd); + cache.cross_keys.resize(config.decoder_layer_configs.size()); + cache.cross_values.resize(config.decoder_layer_configs.size()); + for (size_t layer_idx = 0; layer_idx < config.decoder_layer_configs.size(); + ++layer_idx) { + const LayerConfig& layer_config = config.decoder_layer_configs[layer_idx]; + const Extents2D extents(cache.source_len, + layer_config.kv_heads * layer_config.qkv_dim); + cache.cross_keys[layer_idx] = + MatStorageT("t5_cross_k", extents, allocator, MatPadding::kOdd); + cache.cross_values[layer_idx] = + MatStorageT("t5_cross_v", extents, allocator, MatPadding::kOdd); + } + cache.pad_mask.resize(cache.source_len); + for (size_t i = 0; i < cache.source_len; ++i) { + cache.pad_mask[i] = prompt[i] == kT5GemmaPadId ? 1 : 0; + } +} + +static void AttachT5GemmaEncoderCaches( + const ModelConfig& config, AllQueries& all_queries, + std::vector& encoder_caches, + const Allocator& allocator) { + encoder_caches.resize(all_queries.NumQueries()); + for (size_t qi = 0; qi < all_queries.NumQueries(); ++qi) { + ValidateT5GemmaFreshGeneration(all_queries[qi].initial_pos, + all_queries[qi].prefix_end); + InitT5GemmaEncoderCache(config, all_queries[qi].prompt, encoder_caches[qi], + allocator); + all_queries[qi].t5gemma_encoder_cache = &encoder_caches[qi]; + all_queries[qi].mutable_pos = 0; + all_queries[qi].prev_token = kT5GemmaBosId; + } +} + +static HWY_INLINE size_t T5GemmaEncoderWindowSize(const ModelConfig& config, + size_t layer_idx) { + if (layer_idx < config.encoder_attention_window_sizes.size()) { + return config.encoder_attention_window_sizes[layer_idx]; + } + return config.max_seq_len; +} + +static HWY_INLINE bool T5GemmaEncoderCanAttend(const ModelConfig& config, + const T5GemmaEncoderCache& cache, + size_t layer_idx, + size_t query_pos, + size_t key_pos) { + if (cache.pad_mask[key_pos]) return false; + const size_t window_size = T5GemmaEncoderWindowSize(config, layer_idx); + if (window_size >= config.max_seq_len) return true; + const size_t distance = query_pos > key_pos ? query_pos - key_pos + : key_pos - query_pos; + return distance <= window_size; +} + +static HWY_INLINE float T5GemmaMaybeSoftCap(float score, float cap) { + if (cap == 0.0f) return score; + return cap * tanhf(score / cap); +} + +static void T5GemmaApplyRope(const LayerConfig& layer_config, + const MatPtrT& inv_timescale, float scale, + MatStorageT& q, + MatStorageT& kv, + ThreadingContext& ctx) { + const size_t source_len = q.Rows(); + const size_t qkv_dim = layer_config.qkv_dim; + for (size_t pos = 0; pos < source_len; ++pos) { + for (size_t head = 0; head < layer_config.heads; ++head) { + float* q_row = q.Row(pos) + head * qkv_dim; + if (layer_config.post_qk == PostQKType::HalfRope) { + Rope(q_row, qkv_dim / 2, inv_timescale.PackedScale1(), pos, ctx, + /*worker=*/0); + if (scale != 1.0f) MulByConst(scale, q_row, qkv_dim); + } else { + RopeAndMulBy(scale, q_row, qkv_dim, inv_timescale.PackedScale1(), pos, + ctx, /*worker=*/0); + } + } + for (size_t kv_head = 0; kv_head < layer_config.kv_heads; ++kv_head) { + float* k_row = kv.Row(pos) + kv_head * 2 * qkv_dim; + if (layer_config.post_qk == PostQKType::HalfRope) { + Rope(k_row, qkv_dim / 2, inv_timescale.PackedScale1(), pos, ctx, + /*worker=*/0); + } else { + RopeAndMulBy(/*mul=*/1.0f, k_row, qkv_dim, + inv_timescale.PackedScale1(), pos, ctx, /*worker=*/0); + } + } + } +} + +static void T5GemmaEncoderAttentionReference( + const ModelConfig& config, size_t layer_idx, + const T5GemmaEncoderLayerWeightsPtrs& layer, + T5GemmaEncoderCache& encoder_cache, MatStorageT& pre_att_rms_out, + MatStorageT& q, MatStorageT& kv, + MatStorageT& att_out, MatStorageT& att_sums, + const MatPtrT& inv_timescale, MatMulEnv& env) { + const LayerConfig& layer_config = layer.layer_config; + const size_t source_len = encoder_cache.source_len; + const size_t qkv_dim = layer_config.qkv_dim; + const size_t heads = layer_config.heads; + const size_t kv_heads = layer_config.kv_heads; + const size_t heads_per_kv = heads / kv_heads; + + CallMatMul(pre_att_rms_out, layer.qkv_einsum_w1, /*add=*/nullptr, env, q); + CallMatMul(pre_att_rms_out, layer.qkv_einsum_w2, /*add=*/nullptr, env, kv); + T5GemmaApplyRope(layer_config, inv_timescale, + /*scale=*/1.0f / sqrtf(static_cast(qkv_dim)), q, kv, + env.ctx); + + for (size_t query_pos = 0; query_pos < source_len; ++query_pos) { + for (size_t head = 0; head < heads; ++head) { + const size_t kv_head = head / heads_per_kv; + const float* query = q.Row(query_pos) + head * qkv_dim; + + float max_score = -std::numeric_limits::infinity(); + for (size_t key_pos = 0; key_pos < source_len; ++key_pos) { + if (!T5GemmaEncoderCanAttend(config, encoder_cache, layer_idx, + query_pos, key_pos)) { + continue; + } + const float* key = kv.Row(key_pos) + kv_head * 2 * qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * key[dim]; + } + score = T5GemmaMaybeSoftCap(score, config.att_cap); + max_score = std::max(max_score, score); + } + + float denom = 0.0f; + float* out = att_out.Row(query_pos) + head * qkv_dim; + std::fill(out, out + qkv_dim, 0.0f); + if (max_score == -std::numeric_limits::infinity()) continue; + + for (size_t key_pos = 0; key_pos < source_len; ++key_pos) { + if (!T5GemmaEncoderCanAttend(config, encoder_cache, layer_idx, + query_pos, key_pos)) { + continue; + } + const float* key = kv.Row(key_pos) + kv_head * 2 * qkv_dim; + const float* value = key + qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * key[dim]; + } + score = T5GemmaMaybeSoftCap(score, config.att_cap); + const float weight = expf(score - max_score); + denom += weight; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + out[dim] += weight * value[dim]; + } + } + const float inv_denom = denom == 0.0f ? 0.0f : 1.0f / denom; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + out[dim] *= inv_denom; + } + } + } + + CallMatMul(att_out, layer.att_weights, /*add=*/nullptr, env, att_sums); +} + +static void T5GemmaEncoderFFW(const T5GemmaEncoderLayerWeightsPtrs& layer, + MatStorageT& hidden_states, + MatStorageT& pre_ffw_rms_out, + MatStorageT& c1, MatStorageT& c2, + MatStorageT& ffw_out, MatMulEnv& env) { + RMSNormBatched(hidden_states, layer.pre_ffw_norm_scale, pre_ffw_rms_out, + env.ctx); + +#if GEMMA_FUSED_FFN + const LayerConfig& layer_config = layer.layer_config; + const auto fused = [&](RowPtrsBF C1, IndexRange range_r, IndexRange range_c, + StridedViewBF C2, size_t worker) { + Activation(layer_config.activation, C1, range_r, range_c, C2, env.ctx, + worker); + }; + MMOptions options; + options.SetFunc(fused); + CallTwoMatMul(pre_ffw_rms_out, layer.gating_einsum_w1, + layer.gating_einsum_w2, env, c1, options); +#else + CallMatMul(pre_ffw_rms_out, layer.gating_einsum_w1, /*add=*/nullptr, env, c1); + CallMatMul(pre_ffw_rms_out, layer.gating_einsum_w2, /*add=*/nullptr, env, c2); + ActivationBatched(layer.layer_config.activation, c1, &c2, env.ctx); +#endif + + CallMatMul(c1, layer.linear_w, /*add=*/nullptr, env, ffw_out); + RMSNormInplaceBatched(layer.post_ffw_norm_scale, ffw_out, env.ctx); + AddFromBatched(ffw_out, hidden_states, env.ctx); +} + +static void T5GemmaEncoderLayerReference( + const ModelConfig& config, size_t layer_idx, + const T5GemmaEncoderLayerWeightsPtrs& layer, + T5GemmaEncoderCache& encoder_cache, MatStorageT& pre_att_rms_out, + MatStorageT& q, MatStorageT& kv, + MatStorageT& att_out, MatStorageT& att_sums, + MatStorageT& pre_ffw_rms_out, MatStorageT& c1, + MatStorageT& c2, MatStorageT& ffw_out, + const MatPtrT& inv_timescale, MatMulEnv& env) { + RMSNormBatched(encoder_cache.hidden_states, layer.pre_attention_norm_scale, + pre_att_rms_out, env.ctx); + T5GemmaEncoderAttentionReference(config, layer_idx, layer, encoder_cache, + pre_att_rms_out, q, kv, att_out, att_sums, + inv_timescale, env); + RMSNormInplaceBatched(layer.post_attention_norm_scale, att_sums, env.ctx); + AddFromBatched(att_sums, encoder_cache.hidden_states, env.ctx); + + T5GemmaEncoderFFW(layer, encoder_cache.hidden_states, pre_ffw_rms_out, c1, c2, + ffw_out, env); +} + +static void T5GemmaEncode(const ModelConfig& config, const WeightsPtrs& weights, + const PromptTokens& prompt, + T5GemmaEncoderCache& encoder_cache, + MatMulEnv& env) { + GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(), Zones::kGenEmbed); + HWY_ASSERT(config.is_encoder_decoder); + HWY_ASSERT(encoder_cache.source_len == prompt.size()); + HWY_ASSERT(encoder_cache.hidden_states.Rows() == prompt.size()); + HWY_ASSERT(encoder_cache.hidden_states.Cols() == config.model_dim); + HWY_ASSERT(encoder_cache.pad_mask.size() == prompt.size()); + + for (size_t pos = 0; pos < prompt.size(); ++pos) { + EmbedTokenFromWeights(prompt[pos], pos, config, + weights.t5gemma_encoder_embedding, + encoder_cache.hidden_states); + HWY_DASSERT(encoder_cache.pad_mask[pos] == + (prompt[pos] == kT5GemmaPadId ? 1 : 0)); + } + + const LayerConfig& layer_config = config.encoder_layer_configs[0]; + const size_t source_len = encoder_cache.source_len; + MatStorageT pre_att_rms_out( + MatFactory("t5_e_pre_att", source_len, config.model_dim, + env.ctx.allocator)); + MatStorageT q( + MatFactory("t5_e_q", source_len, + layer_config.heads * layer_config.qkv_dim, + env.ctx.allocator)); + MatStorageT kv( + MatFactory("t5_e_kv", source_len, + 2 * layer_config.kv_heads * layer_config.qkv_dim, + env.ctx.allocator)); + MatStorageT att_out( + MatFactory("t5_e_att_out", source_len, + layer_config.heads * layer_config.qkv_dim, + env.ctx.allocator)); + MatStorageT att_sums( + MatFactory("t5_e_att_sums", source_len, config.model_dim, + env.ctx.allocator)); + MatStorageT pre_ffw_rms_out( + MatFactory("t5_e_pre_ffw", source_len, config.model_dim, + env.ctx.allocator)); + MatStorageT c1( + MatFactory("t5_e_c1", source_len, layer_config.ff_hidden_dim, + env.ctx.allocator)); + MatStorageT c2( + MatFactory("t5_e_c2", source_len, layer_config.ff_hidden_dim, + env.ctx.allocator)); + MatStorageT ffw_out( + MatFactory("t5_e_ffw_out", source_len, config.model_dim, + env.ctx.allocator)); + MatStorageT inv_timescale = + CreateInvTimescale(env.ctx.allocator, layer_config.qkv_dim, + layer_config.post_qk == PostQKType::HalfRope); + q.AllocateAndAttachRowPtrs(env.row_ptrs); + kv.AllocateAndAttachRowPtrs(env.row_ptrs); + att_out.AllocateAndAttachRowPtrs(env.row_ptrs); + att_sums.AllocateAndAttachRowPtrs(env.row_ptrs); + c1.AllocateAndAttachRowPtrs(env.row_ptrs); + c2.AllocateAndAttachRowPtrs(env.row_ptrs); + ffw_out.AllocateAndAttachRowPtrs(env.row_ptrs); + + for (size_t layer_idx = 0; layer_idx < weights.t5gemma_encoder_layers.size(); + ++layer_idx) { + T5GemmaEncoderLayerReference( + config, layer_idx, weights.t5gemma_encoder_layers[layer_idx], + encoder_cache, pre_att_rms_out, q, kv, att_out, att_sums, + pre_ffw_rms_out, c1, c2, ffw_out, inv_timescale, env); + } + RMSNormInplaceBatched(weights.t5gemma_encoder_final_norm_scale, + encoder_cache.hidden_states, env.ctx); +} + +static void T5GemmaPrecomputeCrossAttentionKV( + const WeightsPtrs& weights, T5GemmaEncoderCache& encoder_cache, + MatMulEnv& env) { + HWY_ASSERT(encoder_cache.cross_keys.size() == + weights.t5gemma_decoder_layers.size()); + HWY_ASSERT(encoder_cache.cross_values.size() == + weights.t5gemma_decoder_layers.size()); + for (size_t layer_idx = 0; layer_idx < weights.t5gemma_decoder_layers.size(); + ++layer_idx) { + const T5GemmaDecoderLayerWeightsPtrs& layer = + weights.t5gemma_decoder_layers[layer_idx]; + MatStorageT& cross_k = encoder_cache.cross_keys[layer_idx]; + MatStorageT& cross_v = encoder_cache.cross_values[layer_idx]; + cross_k.AllocateAndAttachRowPtrs(env.row_ptrs); + cross_v.AllocateAndAttachRowPtrs(env.row_ptrs); + CallMatMul(encoder_cache.hidden_states, layer.cross_k_einsum_w, + /*add=*/nullptr, env, cross_k); + CallMatMul(encoder_cache.hidden_states, layer.cross_v_einsum_w, + /*add=*/nullptr, env, cross_v); + } +} + +static void T5GemmaEncodeAllQueries(const ModelConfig& config, + const WeightsPtrs& weights, + AllQueries& all_queries, + MatMulEnv& env) { + for (size_t qi = 0; qi < all_queries.NumQueries(); ++qi) { + T5GemmaEncoderCache* encoder_cache = + all_queries[qi].t5gemma_encoder_cache; + HWY_ASSERT(encoder_cache != nullptr); + T5GemmaEncode(config, weights, all_queries[qi].prompt, *encoder_cache, env); + T5GemmaPrecomputeCrossAttentionKV(weights, *encoder_cache, env); + } +} + +static size_t T5GemmaPromptTokenCount(const AllQueries& all_queries) { + size_t tokens = 0; + for (size_t qi = 0; qi < all_queries.NumQueries(); ++qi) { + tokens += all_queries[qi].prompt.size(); + } + return tokens; +} + +static void T5GemmaEmbedDecoderTokens(const ModelConfig& config, + const WeightsPtrs& weights, + Activations& activations, + QBatch& qbatch, + MatMulEnv& env) { + activations.SetBatchSize(qbatch.Size()); + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + EmbedTokenFromWeights(qbatch.PrevToken(qi), qi, config, + weights.t5gemma_decoder_embedding, activations.x); + } +} + +static HWY_INLINE size_t T5GemmaDecoderWindowSize(const ModelConfig& config, + size_t layer_idx) { + if (layer_idx < config.decoder_attention_window_sizes.size()) { + return config.decoder_attention_window_sizes[layer_idx]; + } + return config.max_seq_len; +} + +static HWY_INLINE size_t T5GemmaDecoderStartPos(const ModelConfig& config, + size_t layer_idx, + size_t query_pos) { + const size_t window_size = T5GemmaDecoderWindowSize(config, layer_idx); + if (window_size >= config.max_seq_len || query_pos < window_size) return 0; + return query_pos + 1 - window_size; +} + +static void T5GemmaApplyDecoderRope(const LayerConfig& layer_config, + const MatPtrT& inv_timescale, + float scale, MatPtrT& q, + MatStorageT& kv, QBatch& qbatch, + ThreadingContext& ctx) { + const size_t qkv_dim = layer_config.qkv_dim; + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + const size_t pos = qbatch.Pos(qi); + for (size_t head = 0; head < layer_config.heads; ++head) { + float* q_row = q.Row(qi) + head * qkv_dim; + if (layer_config.post_qk == PostQKType::HalfRope) { + Rope(q_row, qkv_dim / 2, inv_timescale.PackedScale1(), pos, ctx, + /*worker=*/0); + if (scale != 1.0f) MulByConst(scale, q_row, qkv_dim); + } else { + RopeAndMulBy(scale, q_row, qkv_dim, inv_timescale.PackedScale1(), pos, + ctx, /*worker=*/0); + } + } + for (size_t kv_head = 0; kv_head < layer_config.kv_heads; ++kv_head) { + float* k_row = kv.Row(qi) + kv_head * 2 * qkv_dim; + if (layer_config.post_qk == PostQKType::HalfRope) { + Rope(k_row, qkv_dim / 2, inv_timescale.PackedScale1(), pos, ctx, + /*worker=*/0); + } else { + RopeAndMulBy(/*mul=*/1.0f, k_row, qkv_dim, + inv_timescale.PackedScale1(), pos, ctx, /*worker=*/0); + } + } + } +} + +static void T5GemmaWriteDecoderKV(const T5GemmaDecoderLayerWeightsPtrs& layer, + size_t layer_idx, const MatStorageT& kv, + QBatch& qbatch) { + const LayerConfig& layer_config = layer.layer_config; + const size_t cache_layer_size = layer_config.CacheLayerSize(); + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + if (qbatch.KV(qi).IsTiled()) { + HWY_ABORT( + "T5Gemma reference decoder self-attention currently requires the " + "plain KV cache; tiled/transposed KV cache support is not wired yet."); + } + const size_t pos = qbatch.Pos(qi); + HWY_ASSERT(pos < qbatch.KV(qi).SeqLen()); + KV_t* dst = + qbatch.KV(qi).kv_cache.Row(pos) + layer_idx * cache_layer_size; + const float* src = kv.Row(qi); + for (size_t i = 0; i < cache_layer_size; ++i) { + dst[i] = hwy::ConvertScalarTo(src[i]); + } + } +} + +static void T5GemmaDecoderSelfAttentionReference( + const ModelConfig& config, size_t layer_idx, + const T5GemmaDecoderLayerWeightsPtrs& layer, Activations& activations, + MatStorageT& kv, const MatPtrT& inv_timescale, + QBatch& qbatch, MatMulEnv& env) { + const LayerConfig& layer_config = layer.layer_config; + const size_t qkv_dim = layer_config.qkv_dim; + const size_t heads = layer_config.heads; + const size_t kv_heads = layer_config.kv_heads; + const size_t heads_per_kv = heads / kv_heads; + const size_t cache_layer_size = layer_config.CacheLayerSize(); + + RMSNormBatched(activations.x, layer.pre_self_attention_norm_scale, + activations.attention.pre_att_rms_out, env.ctx); + CallMatMul(activations.attention.pre_att_rms_out, layer.self_qkv_einsum_w1, + /*add=*/nullptr, env, activations.attention.q); + CallMatMul(activations.attention.pre_att_rms_out, layer.self_qkv_einsum_w2, + /*add=*/nullptr, env, kv); + T5GemmaApplyDecoderRope(layer_config, inv_timescale, + /*scale=*/1.0f / sqrtf(static_cast(qkv_dim)), + activations.attention.q, kv, qbatch, env.ctx); + T5GemmaWriteDecoderKV(layer, layer_idx, kv, qbatch); + + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + const size_t query_pos = qbatch.Pos(qi); + const size_t start_pos = + T5GemmaDecoderStartPos(config, layer_idx, query_pos); + for (size_t head = 0; head < heads; ++head) { + const size_t kv_head = head / heads_per_kv; + const float* query = activations.attention.q.Row(qi) + head * qkv_dim; + + float max_score = -std::numeric_limits::infinity(); + for (size_t key_pos = start_pos; key_pos <= query_pos; ++key_pos) { + const KV_t* key = + qbatch.KV(qi).kv_cache.Row(key_pos) + layer_idx * cache_layer_size + + kv_head * 2 * qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * hwy::ConvertScalarTo(key[dim]); + } + score = T5GemmaMaybeSoftCap(score, config.att_cap); + max_score = std::max(max_score, score); + } + + float denom = 0.0f; + float* out = activations.attention.att_out.Row(qi) + head * qkv_dim; + std::fill(out, out + qkv_dim, 0.0f); + for (size_t key_pos = start_pos; key_pos <= query_pos; ++key_pos) { + const KV_t* key = + qbatch.KV(qi).kv_cache.Row(key_pos) + layer_idx * cache_layer_size + + kv_head * 2 * qkv_dim; + const KV_t* value = key + qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * hwy::ConvertScalarTo(key[dim]); + } + score = T5GemmaMaybeSoftCap(score, config.att_cap); + const float weight = expf(score - max_score); + denom += weight; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + out[dim] += weight * hwy::ConvertScalarTo(value[dim]); + } + } + const float inv_denom = denom == 0.0f ? 0.0f : 1.0f / denom; + for (size_t dim = 0; dim < qkv_dim; ++dim) out[dim] *= inv_denom; + } + } + + CallMatMul(activations.attention.att_out, layer.self_att_weights, + /*add=*/nullptr, env, activations.attention.att_sums); + RMSNormInplaceBatched(layer.post_self_attention_norm_scale, + activations.attention.att_sums, env.ctx); + AddFromBatched(activations.attention.att_sums, activations.x, env.ctx); +} + +static void T5GemmaDecoderCrossAttentionReference( + const ModelConfig& config, size_t layer_idx, + const T5GemmaDecoderLayerWeightsPtrs& layer, Activations& activations, + QBatch& qbatch, MatMulEnv& env) { + const LayerConfig& layer_config = layer.layer_config; + const size_t qkv_dim = layer_config.qkv_dim; + const size_t heads = layer_config.heads; + const size_t kv_heads = layer_config.kv_heads; + const size_t heads_per_kv = heads / kv_heads; + + RMSNormBatched(activations.x, layer.pre_cross_attention_norm_scale, + activations.attention.pre_att_rms_out, env.ctx); + CallMatMul(activations.attention.pre_att_rms_out, layer.cross_q_einsum_w, + /*add=*/nullptr, env, activations.attention.q); + + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + const T5GemmaEncoderCache* encoder_cache = qbatch.T5EncoderCache(qi); + HWY_ASSERT(encoder_cache != nullptr); + HWY_ASSERT(layer_idx < encoder_cache->cross_keys.size()); + HWY_ASSERT(layer_idx < encoder_cache->cross_values.size()); + const size_t source_len = encoder_cache->source_len; + const MatStorageT& cross_k = encoder_cache->cross_keys[layer_idx]; + const MatStorageT& cross_v = encoder_cache->cross_values[layer_idx]; + HWY_ASSERT(cross_k.Rows() == source_len); + HWY_ASSERT(cross_v.Rows() == source_len); + HWY_ASSERT(cross_k.Cols() == kv_heads * qkv_dim); + HWY_ASSERT(cross_v.Cols() == kv_heads * qkv_dim); + + for (size_t head = 0; head < heads; ++head) { + const size_t kv_head = head / heads_per_kv; + const float* query = activations.attention.q.Row(qi) + head * qkv_dim; + + float max_score = -std::numeric_limits::infinity(); + for (size_t source_pos = 0; source_pos < source_len; ++source_pos) { + if (encoder_cache->pad_mask[source_pos]) continue; + const float* key = cross_k.Row(source_pos) + kv_head * qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * key[dim]; + } + score = T5GemmaMaybeSoftCap(score / sqrtf(static_cast(qkv_dim)), + config.att_cap); + max_score = std::max(max_score, score); + } + + float denom = 0.0f; + float* out = activations.attention.att_out.Row(qi) + head * qkv_dim; + std::fill(out, out + qkv_dim, 0.0f); + if (max_score == -std::numeric_limits::infinity()) continue; + + for (size_t source_pos = 0; source_pos < source_len; ++source_pos) { + if (encoder_cache->pad_mask[source_pos]) continue; + const float* key = cross_k.Row(source_pos) + kv_head * qkv_dim; + const float* value = cross_v.Row(source_pos) + kv_head * qkv_dim; + float score = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + score += query[dim] * key[dim]; + } + score = T5GemmaMaybeSoftCap(score / sqrtf(static_cast(qkv_dim)), + config.att_cap); + const float weight = expf(score - max_score); + denom += weight; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + out[dim] += weight * value[dim]; + } + } + const float inv_denom = denom == 0.0f ? 0.0f : 1.0f / denom; + for (size_t dim = 0; dim < qkv_dim; ++dim) out[dim] *= inv_denom; + } + } + + CallMatMul(activations.attention.att_out, layer.cross_att_weights, + /*add=*/nullptr, env, activations.attention.att_sums); + RMSNormInplaceBatched(layer.post_cross_attention_norm_scale, + activations.attention.att_sums, env.ctx); + AddFromBatched(activations.attention.att_sums, activations.x, env.ctx); +} + +static void T5GemmaDecoderFFW(const T5GemmaDecoderLayerWeightsPtrs& layer, + Activations& activations, MatMulEnv& env) { + RMSNormBatched(activations.x, layer.pre_ffw_norm_scale, + activations.pre_ffw_rms_out, env.ctx); + +#if GEMMA_FUSED_FFN + const LayerConfig& layer_config = layer.layer_config; + const auto fused = [&](RowPtrsBF C1, IndexRange range_r, IndexRange range_c, + StridedViewBF C2, size_t worker) { + Activation(layer_config.activation, C1, range_r, range_c, C2, env.ctx, + worker); + }; + MMOptions options; + options.SetFunc(fused); + CallTwoMatMul(activations.pre_ffw_rms_out, layer.gating_einsum_w1, + layer.gating_einsum_w2, env, activations.C1, options); +#else + CallMatMul(activations.pre_ffw_rms_out, layer.gating_einsum_w1, + /*add=*/nullptr, env, activations.C1); + CallMatMul(activations.pre_ffw_rms_out, layer.gating_einsum_w2, + /*add=*/nullptr, env, activations.C2); + ActivationBatched(layer.layer_config.activation, activations.C1, + &activations.C2, env.ctx); +#endif + + CallMatMul(activations.C1, layer.linear_w, /*add=*/nullptr, env, + activations.ffw_out); + RMSNormInplaceBatched(layer.post_ffw_norm_scale, activations.ffw_out, + env.ctx); + AddFromBatched(activations.ffw_out, activations.x, env.ctx); +} + +static void T5GemmaComputeLogitsChunked(const WeightsPtrs& weights, + Activations& activations, + MatMulEnv& env, + uint8_t** chunk_row_ptrs, + const hwy::BitSet4096<>* non_eos, + int* greedy_tokens, + float* greedy_logits) { + constexpr size_t kLogitsChunk = kMaxNC; + const size_t vocab_size = weights.t5gemma_decoder_embedding.Rows(); + const size_t model_dim = weights.t5gemma_decoder_embedding.Cols(); + for (size_t start = 0; start < vocab_size; start += kLogitsChunk) { + const size_t rows = std::min(kLogitsChunk, vocab_size - start); + MatPtr embedding_chunk("dec_emb_chunk", + weights.t5gemma_decoder_embedding.GetType(), + Extents2D(rows, model_dim)); + embedding_chunk.SetScale(weights.t5gemma_decoder_embedding.Scale()); + embedding_chunk.SetPtr( + const_cast( + weights.t5gemma_decoder_embedding.RowBytes(start)), + weights.t5gemma_decoder_embedding.Stride()); + + MatPtrT logits_chunk( + "logits_chunk", Extents2D(activations.logits.Rows(), rows)); + logits_chunk.SetPtr(activations.logits.Row(0) + start, + activations.logits.Stride()); + for (size_t qi = 0; qi < activations.logits.Rows(); ++qi) { + chunk_row_ptrs[qi] = + reinterpret_cast(activations.logits.Row(qi) + start); + } + logits_chunk.AttachRowPtrs(chunk_row_ptrs); + + CallMatMul(activations.x_bf, embedding_chunk, /*add=*/nullptr, env, + logits_chunk); + + if (greedy_tokens != nullptr) { + for (size_t qi = 0; qi < activations.logits.Rows(); ++qi) { + if (non_eos != nullptr && !non_eos->Get(qi)) continue; + const TokenAndProb chunk_best = + ArgmaxAndMax(Logits(activations.logits.Row(qi) + start, rows)); + if (chunk_best.prob > greedy_logits[qi]) { + greedy_logits[qi] = chunk_best.prob; + greedy_tokens[qi] = static_cast(start + chunk_best.token); + } + } + } + } +} + +static void T5GemmaDecoderTransformer(const ModelConfig& config, + const WeightsPtrs& weights, + Activations& activations, + MatStorageT& decoder_kv, + const MatPtrT& inv_timescale, + QBatch& qbatch, MatMulEnv& env) { + for (size_t layer_idx = 0; layer_idx < weights.t5gemma_decoder_layers.size(); + ++layer_idx) { + const T5GemmaDecoderLayerWeightsPtrs& layer = + weights.t5gemma_decoder_layers[layer_idx]; + T5GemmaDecoderSelfAttentionReference(config, layer_idx, layer, activations, + decoder_kv, inv_timescale, qbatch, + env); + T5GemmaDecoderCrossAttentionReference(config, layer_idx, layer, + activations, qbatch, env); + T5GemmaDecoderFFW(layer, activations, env); + } +} + +static bool T5GemmaUseGreedyFastPath(const RuntimeConfig& runtime_config) { + return !runtime_config.sample_func && !runtime_config.accept_token && + runtime_config.top_k == 1 && runtime_config.temperature == 0.0f; +} + +static void T5GemmaSampleAndStream(const ModelConfig& config, + const RuntimeConfig& runtime_config, + const WeightsPtrs& weights, + const SampleFunc& sample_token, + Activations& activations, QBatch& qbatch, + MatMulEnv& env, uint8_t** chunk_row_ptrs, + std::vector& greedy_tokens, + std::vector& greedy_logits, + hwy::BitSet4096<>& non_eos, + TimingInfo& timing_info) { + HWY_DASSERT(qbatch.Size() == activations.x.Rows()); + + RMSNormBatched(activations.x, weights.t5gemma_decoder_final_norm_scale, + activations.x_bf, env.ctx); + const bool greedy_fast_path = T5GemmaUseGreedyFastPath(runtime_config); + if (greedy_fast_path) { + std::fill(greedy_tokens.begin(), greedy_tokens.end(), 0); + std::fill(greedy_logits.begin(), greedy_logits.end(), + -std::numeric_limits::infinity()); + } + { + GCPP_ZONE(env.ctx, /*worker=*/0, Zones::kGenEmbeddingMatmul); + T5GemmaComputeLogitsChunked( + weights, activations, env, chunk_row_ptrs, + greedy_fast_path ? &non_eos : nullptr, + greedy_fast_path ? greedy_tokens.data() : nullptr, + greedy_fast_path ? greedy_logits.data() : nullptr); + } + if (!greedy_fast_path) { + MaybeLogitsSoftCapBatched(config.final_cap, activations.logits, non_eos, + env.ctx); + } + + timing_info.NotifyGenerated(non_eos.Count()); + ParallelFor( + Parallelism::kFlat, qbatch.Size(), env.ctx, + /*cluster_idx=*/0, Callers::kSampleAndStream, + [&](size_t qi, size_t worker) { + if (!non_eos.Get(qi)) return; + + const size_t pos = qbatch.Pos(qi); + TokenAndProb tp; + if (greedy_fast_path) { + tp.token = greedy_tokens[qi]; + tp.prob = 1.0f; + } else { + tp = sample_token(qi, pos, activations.logits.RowSpan(qi), worker); + } + activations.sampled.Row(qi)[0] = static_cast(pos); + activations.sampled.Row(qi)[1] = static_cast(tp.token); + activations.sampled.Row(qi)[2] = hwy::BitCastScalar(tp.prob); + }); + + non_eos.Foreach([&](size_t qi) { + const size_t pos = activations.sampled.Row(qi)[0]; + const int token = static_cast(activations.sampled.Row(qi)[1]); + const float prob = + hwy::BitCastScalar(activations.sampled.Row(qi)[2]); + StreamAndUpdateEOS(qi, pos, token, prob, config, runtime_config, qbatch, + /*update_pos=*/true, non_eos); + }); +} + +static void T5GemmaGenerateT(const ModelConfig& config, + const RuntimeConfig& runtime_config, + const AesCtrEngine& engine, + const WeightsPtrs& weights, + Activations& activations, QBatch& qbatch, + MatMulEnv& env, TimingInfo& timing_info) { + hwy::BitSet4096<> non_eos; + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + non_eos.Set(qi); + qbatch.PrevToken(qi) = kT5GemmaBosId; + } + + const SampleFunc sample_token = + ChooseSampleFunc(runtime_config, engine, env.ctx); + const size_t max_gen_steps = + HWY_MIN(runtime_config.max_generated_tokens, qbatch.KV(0).SeqLen()); + const LayerConfig& layer_config = config.decoder_layer_configs[0]; + MatStorageT decoder_kv( + MatFactory("t5_d_kv", qbatch.Size(), + 2 * layer_config.kv_heads * layer_config.qkv_dim, + env.ctx.allocator)); + decoder_kv.AllocateAndAttachRowPtrs(env.row_ptrs); + MatStorageT inv_timescale = + CreateInvTimescale(env.ctx.allocator, layer_config.qkv_dim, + layer_config.post_qk == PostQKType::HalfRope); + auto logits_chunk_row_ptrs = hwy::AllocateAligned(qbatch.Size()); + std::vector greedy_tokens(qbatch.Size()); + std::vector greedy_logits(qbatch.Size()); + timing_info.generate_start = hwy::platform::Now(); + for (size_t gen = 0; gen < max_gen_steps && non_eos.Any(); ++gen) { + T5GemmaEmbedDecoderTokens(config, weights, activations, qbatch, env); + T5GemmaDecoderTransformer(config, weights, activations, decoder_kv, + inv_timescale, qbatch, env); + T5GemmaSampleAndStream(config, runtime_config, weights, sample_token, + activations, qbatch, env, + logits_chunk_row_ptrs.get(), greedy_tokens, + greedy_logits, non_eos, timing_info); + } + timing_info.NotifyGenerateDone(); +} + // `x_row` indicates which row of `x` to write to. // `pos` is the *token*'s position for `AddAbsolutePositionalEmbeddings`, not // the start of the batch, because this is called for batches of tokens in @@ -179,28 +1011,11 @@ EmbedMMToken(int token, size_t x_row, size_t pos, size_t pos_in_prompt, return image_token_position; } - const size_t model_dim = model_config.model_dim; - const float emb_scaling = EmbeddingScaling(model_dim); - - HWY_DASSERT(token >= 0); - HWY_DASSERT(token < static_cast(model_config.vocab_size)); - - CallUpcasted(&weights.embedder_input_embedding, [&](const auto* weights_t) { - // Using `Stride` to compute the offset works for both NUQ (because we use - // an offset and NUQ is never padded) and padded, because non-NUQ types are - // seekable, hence the offset can also skip any padding. - const size_t embedding_ofs = token * weights_t->Stride(); - HWY_ASSERT(weights_t->Cols() == model_dim); - const auto embedding_span = - MakeSpan(weights_t->Row(0), embedding_ofs + model_dim); - const hn::ScalableTag df; - DecompressAndZeroPad(df, embedding_span, embedding_ofs, x.Row(x_row), - model_dim); - MulByConst(emb_scaling * weights_t->Scale(), x.Row(x_row), model_dim); - }); + EmbedTokenFromWeights(token, x_row, model_config, + weights.embedder_input_embedding, x); if (model_config.absolute_pe) { - AddAbsolutePositionalEmbeddings(x.Row(x_row), model_dim, pos); + AddAbsolutePositionalEmbeddings(x.Row(x_row), model_config.model_dim, pos); } return image_token_position; } @@ -705,6 +1520,23 @@ void GenerateSingleT(const PromptTokens& prompt, size_t pos, size_t prefix_end, const AesCtrEngine& engine, const WeightsPtrs& weights, KVCache& kv_cache, MatMulEnv& env, TimingInfo& timing_info) { + if (config.is_encoder_decoder) { + ValidateT5GemmaFreshGeneration(pos, prefix_end); + std::vector encoder_caches; + AllQueries all_queries(prompt, pos, prefix_end, + hwy::Span(&kv_cache, 1)); + AttachT5GemmaEncoderCaches(config, all_queries, encoder_caches, + env.ctx.allocator); + timing_info.prefill_start = hwy::platform::Now(); + T5GemmaEncodeAllQueries(config, weights, all_queries, env); + timing_info.NotifyPrefill(T5GemmaPromptTokenCount(all_queries)); + Activations activations(runtime_config, config, /*batch_size=*/1, + kv_cache.SeqLen(), env.ctx, env.row_ptrs); + QBatch qbatch(/*start=*/0, /*max_size=*/1, all_queries); + T5GemmaGenerateT(config, runtime_config, engine, weights, activations, + qbatch, env, timing_info); + return; + } Activations activations(runtime_config, config, runtime_config.prefill_tbatch_size, kv_cache.SeqLen(), env.ctx, env.row_ptrs); @@ -723,6 +1555,23 @@ void GenerateBatchT(const ModelConfig& config, const AesCtrEngine& engine, const WeightsPtrs& weights, AllQueries& all_queries, MatMulEnv& env, TimingInfo& timing_info) { + if (config.is_encoder_decoder) { + std::vector encoder_caches; + AttachT5GemmaEncoderCaches(config, all_queries, encoder_caches, + env.ctx.allocator); + timing_info.prefill_start = hwy::platform::Now(); + T5GemmaEncodeAllQueries(config, weights, all_queries, env); + timing_info.NotifyPrefill(T5GemmaPromptTokenCount(all_queries)); + const size_t max_batch_size = HWY_MAX(runtime_config.decode_qbatch_size, + runtime_config.prefill_tbatch_size); + Activations activations(runtime_config, config, max_batch_size, + all_queries[0].kv_cache.SeqLen(), env.ctx, + env.row_ptrs); + QBatch qbatch(/*start=*/0, runtime_config.decode_qbatch_size, all_queries); + T5GemmaGenerateT(config, runtime_config, engine, weights, activations, + qbatch, env, timing_info); + return; + } const size_t max_batch_size = HWY_MAX(runtime_config.decode_qbatch_size, runtime_config.prefill_tbatch_size); Activations activations(runtime_config, config, max_batch_size, diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index b7204aee..f25e0586 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -41,6 +41,18 @@ static size_t CappedSeqLen(const ModelConfig& config, return inference_args.seq_len; } +static const std::vector& KVLayerConfigs( + const ModelConfig& config) { + return config.is_encoder_decoder ? config.decoder_layer_configs + : config.layer_configs; +} + +static const std::vector& KVAttentionWindowSizes( + const ModelConfig& config) { + return config.is_encoder_decoder ? config.decoder_attention_window_sizes + : config.attention_window_sizes; +} + KVCache::KVCache(const Extents2D& kv_extents, size_t num_layers, size_t kv_heads, size_t qkv_dim, const Allocator& allocator) : num_layers(num_layers), @@ -84,15 +96,18 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const Allocator& allocator) : KVCache( Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()), - config.layer_configs.size(), config.layer_configs[0].kv_heads, - config.layer_configs[0].qkv_dim, allocator) {} + KVLayerConfigs(config).size(), KVLayerConfigs(config)[0].kv_heads, + KVLayerConfigs(config)[0].qkv_dim, allocator) {} KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const RuntimeConfig& runtime_config, const Allocator& allocator) : allocator_(allocator) { + const std::vector& kv_layer_configs = KVLayerConfigs(config); + const std::vector& kv_attention_window_sizes = + KVAttentionWindowSizes(config); - num_layers = config.num_layers; + num_layers = kv_layer_configs.size(); // 1. Build non-uniform offset tables dynamically layer_flat_offsets.resize(num_layers, 0); @@ -106,22 +121,23 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, for (size_t i = 0; i < num_layers; ++i) { layer_flat_offsets[i] = static_cast(flat_accum); - flat_accum += config.layer_configs[i].CacheLayerSize(); + flat_accum += kv_layer_configs[i].CacheLayerSize(); layer_k_v_offsets[i] = static_cast(k_v_accum); - size_t rounded_dim = hwy::RoundUpTo(config.layer_configs[i].qkv_dim, kMaxBF16PerVector); + size_t rounded_dim = + hwy::RoundUpTo(kv_layer_configs[i].qkv_dim, kMaxBF16PerVector); rounded_qkv_dims[i] = static_cast(rounded_dim); - k_v_accum += config.layer_configs[i].kv_heads * rounded_dim; + k_v_accum += kv_layer_configs[i].kv_heads * rounded_dim; - max_qkv_dim = HWY_MAX(max_qkv_dim, config.layer_configs[i].qkv_dim); - max_kv_heads = HWY_MAX(max_kv_heads, config.layer_configs[i].kv_heads); + max_qkv_dim = HWY_MAX(max_qkv_dim, kv_layer_configs[i].qkv_dim); + max_kv_heads = HWY_MAX(max_kv_heads, kv_layer_configs[i].kv_heads); } k_v_cols = static_cast(k_v_accum); // Since we also store legacy homogeneous variables (used by tests/old code), // we default them to Layer 0 values. - kv_heads = config.layer_configs[0].kv_heads; - qkv_dim = config.layer_configs[0].qkv_dim; + kv_heads = kv_layer_configs[0].kv_heads; + qkv_dim = kv_layer_configs[0].qkv_dim; rounded_qkv_dim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); // clang-format off @@ -190,9 +206,10 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t total_num_tiles = 0; for (size_t i = 0; i < num_layers; ++i) { total_num_tiles += - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, + num_tiles_per_head(kv_attention_window_sizes[i], + runtime_config.prefill_tbatch_size, config.max_seq_len) * - config.layer_configs[i].kv_heads; + kv_layer_configs[i].kv_heads; } Extents2D extents(total_num_tiles, max_tile_length); compact_kv_cache_ptr = MatPtr("kv_tiled", kv_cache_type, extents); @@ -209,13 +226,14 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_head_ptrs.clear(); kv_head_ptrs.reserve(num_layers * max_kv_heads); for (size_t i = 0; i < num_layers; ++i) { - size_t layer_tile_length = 2 * config.layer_configs[i].qkv_dim * kTileSize; + size_t layer_tile_length = 2 * kv_layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { layer_tile_length += 2 * sizeof(BF16) * kTileSize; } - for (size_t kv = 0; kv < config.layer_configs[i].kv_heads; ++kv) { + for (size_t kv = 0; kv < kv_layer_configs[i].kv_heads; ++kv) { size_t num_tiles_per_kv_head = - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, + num_tiles_per_head(kv_attention_window_sizes[i], + runtime_config.prefill_tbatch_size, config.max_seq_len); MatPtr kv_ptr("kv_ptr", kv_cache_type, Extents2D(num_tiles_per_kv_head, layer_tile_length)); diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc index 9be75d94..17ef7b92 100644 --- a/gemma/kv_cache_test.cc +++ b/gemma/kv_cache_test.cc @@ -45,5 +45,23 @@ TEST(KVCacheTest, KVCacheToPtrs) { } } +TEST(KVCacheTest, EncoderDecoderUsesDecoderLayerConfig) { + ModelConfig model_config(Model::T5GEMMA_S_S, Type::kSFP, + PromptWrapping::GEMMA_PT); + ASSERT_TRUE(model_config.is_encoder_decoder); + ASSERT_FALSE(model_config.decoder_layer_configs.empty()); + InferenceArgs inference_args; + inference_args.seq_len = 128; + ThreadingArgs threading_args; + ThreadingContext ctx(threading_args); + + KVCache cache(model_config, inference_args, ctx.allocator); + + EXPECT_EQ(cache.num_layers, model_config.decoder_layer_configs.size()); + EXPECT_EQ(cache.kv_heads, model_config.decoder_layer_configs[0].kv_heads); + EXPECT_EQ(cache.qkv_dim, model_config.decoder_layer_configs[0].qkv_dim); + EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols()); +} + } // namespace } // namespace gcpp diff --git a/gemma/model_store.cc b/gemma/model_store.cc index 67aa3068..09d29449 100644 --- a/gemma/model_store.cc +++ b/gemma/model_store.cc @@ -224,6 +224,8 @@ static int DeduceLayerTypes(const BlobReader& reader) { int layer_types = 0; bool has_key_norm = false; bool has_query_norm = false; + bool has_t5gemma_encoder = false; + bool has_t5gemma_decoder = false; for (size_t key_idx = 0; key_idx < reader.Keys().size(); ++key_idx) { const std::string& key = reader.Keys()[key_idx]; if (key.find("qkv_ein_w") != std::string::npos) { // NOLINT @@ -241,10 +243,21 @@ static int DeduceLayerTypes(const BlobReader& reader) { if (key.find("query_norm") != std::string::npos) { // NOLINT has_query_norm = true; } + if (key.find("enc_embedding") != std::string::npos || + key.find("e_qkv_") != std::string::npos) { // NOLINT + has_t5gemma_encoder = true; + } + if (key.find("dec_embedding") != std::string::npos || + key.find("d_qkv_") != std::string::npos) { // NOLINT + has_t5gemma_decoder = true; + } } if (has_key_norm && has_query_norm) { layer_types |= kDeducedKqNorm; } + if (has_t5gemma_encoder && has_t5gemma_decoder) { + layer_types |= kDeducedT5Gemma; + } return layer_types; } diff --git a/gemma/query.h b/gemma/query.h index 36e8ee5c..88a6a72e 100644 --- a/gemma/query.h +++ b/gemma/query.h @@ -16,6 +16,9 @@ #ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_QUERY_H_ #define THIRD_PARTY_GEMMA_CPP_GEMMA_QUERY_H_ +#include +#include + #include #include "gemma/gemma_args.h" @@ -26,6 +29,14 @@ namespace gcpp { +struct T5GemmaEncoderCache { + MatStorageT hidden_states; + std::vector> cross_keys; + std::vector> cross_values; + std::vector pad_mask; + size_t source_len = 0; +}; + struct PerQuery { PromptTokens prompt; @@ -41,6 +52,10 @@ struct PerQuery { KVCachePtr kv_cache; + // Non-owning pointer to encoder outputs for encoder-decoder models. The + // owner lives alongside the query batch during generation. + T5GemmaEncoderCache* t5gemma_encoder_cache = nullptr; + // Previous token generated for this query, or the last prompt token. Will be // fed into the next Transformer() call. int prev_token = 0; @@ -161,6 +176,9 @@ class QBatch { return queries_[QueryIdx(qi)].prefix_end; } KVCachePtr& KV(size_t qi) const { return queries_[QueryIdx(qi)].kv_cache; } + T5GemmaEncoderCache*& T5EncoderCache(size_t qi) const { + return queries_[QueryIdx(qi)].t5gemma_encoder_cache; + } int& PrevToken(size_t qi) { return queries_[QueryIdx(qi)].prev_token; } // let query_idx_[to] point to the from in the queries_; this is only used if diff --git a/gemma/run.cc b/gemma/run.cc index bc1877a4..5a3ecd8d 100644 --- a/gemma/run.cc +++ b/gemma/run.cc @@ -208,6 +208,11 @@ void ReplGemma(const GemmaArgs& args, const Gemma& gemma, KVCache& kv_cache, config.wrapping, abs_pos, prompt_string); prompt_size = prompt.size(); } + if (config.is_encoder_decoder) { + // Encoder-decoder models consume the source prompt in the encoder, then + // stream only decoder outputs. + prompt_size = 0; + } if constexpr (kVerboseLogTokens) { for (int i = 0; i < prompt_size; ++i) { diff --git a/gemma/tensor_info.cc b/gemma/tensor_info.cc index 0d48b495..ed40b2ef 100644 --- a/gemma/tensor_info.cc +++ b/gemma/tensor_info.cc @@ -123,6 +123,38 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) { }); } +void TensorInfoRegistry::AddT5GemmaModelTensors(const ModelConfig& config) { + const std::string no_suffix; + Add(no_suffix, { + .base_name = "enc_embedding", + .source_names = {"model.encoder.embed_tokens.weight"}, + .axes = {0, 1}, + .shape = {config.vocab_size, config.model_dim}, + .min_size = Type::kBF16, + }); + Add(no_suffix, { + .base_name = "dec_embedding", + .source_names = {"model.decoder.embed_tokens.weight"}, + .axes = {0, 1}, + .shape = {config.vocab_size, config.model_dim}, + .min_size = Type::kBF16, + }); + Add(no_suffix, { + .base_name = "enc_final_norm", + .source_names = {"model.encoder.norm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(no_suffix, { + .base_name = "dec_final_norm", + .source_names = {"model.decoder.norm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); +} + // Returns the tensors for the given image layer config. void TensorInfoRegistry::AddImageLayerTensors(const ModelConfig& config, const LayerConfig& layer_config, @@ -293,6 +325,235 @@ void TensorInfoRegistry::AddImageLayerTensors(const ModelConfig& config, }); } +void TensorInfoRegistry::AddT5GemmaEncoderLayerTensors( + const ModelConfig& config, const LayerConfig& layer_config, + const size_t layer_idx) { + const std::string suffix = LayerSuffix(layer_idx); + Add(suffix, { + .base_name = "e_qkv", + .source_names = {"model.encoder.layers.self_attn.qkv"}, + .axes = {0, 1, 2}, + .shape = {layer_config.heads + 2 * layer_config.kv_heads, + layer_config.qkv_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "e_qkv1", + .shape = {layer_config.heads * layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "e_qkv2", + .shape = {2 * layer_config.kv_heads * layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "e_att", + .source_names = {"model.encoder.layers.self_attn.o_proj"}, + .preshape = {layer_config.heads, layer_config.qkv_dim, + config.model_dim}, + .axes = {0, 2, 1}, + .shape = {layer_config.heads, config.model_dim, + layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "e_att_w", + .shape = {config.model_dim, + layer_config.heads * layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "e_gate", + .source_names = {"model.encoder.layers.mlp.gating_einsum"}, + .axes = {0, 1, 2}, + .shape = {2, layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "e_gate1", + .shape = {layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "e_gate2", + .shape = {layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "e_lin", + .source_names = {"model.encoder.layers.mlp.down_proj.weight"}, + .axes = {0, 1}, + .shape = {config.model_dim, layer_config.ff_hidden_dim}, + }); + Add(suffix, { + .base_name = "e_pre_att", + .source_names = { + "model.encoder.layers.pre_self_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "e_post_att", + .source_names = { + "model.encoder.layers.post_self_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "e_pre_ff", + .source_names = { + "model.encoder.layers.pre_feedforward_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "e_post_ff", + .source_names = { + "model.encoder.layers.post_feedforward_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); +} + +void TensorInfoRegistry::AddT5GemmaDecoderLayerTensors( + const ModelConfig& config, const LayerConfig& layer_config, + const size_t layer_idx) { + const std::string suffix = LayerSuffix(layer_idx); + Add(suffix, { + .base_name = "d_qkv", + .source_names = {"model.decoder.layers.self_attn.qkv"}, + .axes = {0, 1, 2}, + .shape = {layer_config.heads + 2 * layer_config.kv_heads, + layer_config.qkv_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "d_qkv1", + .shape = {layer_config.heads * layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "d_qkv2", + .shape = {2 * layer_config.kv_heads * layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "d_att", + .source_names = {"model.decoder.layers.self_attn.o_proj"}, + .preshape = {layer_config.heads, layer_config.qkv_dim, + config.model_dim}, + .axes = {0, 2, 1}, + .shape = {layer_config.heads, config.model_dim, + layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "d_att_w", + .shape = {config.model_dim, + layer_config.heads * layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "dc_q", + .source_names = {"model.decoder.layers.cross_attn.q_proj"}, + .axes = {0, 1, 2}, + .shape = {layer_config.heads, layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "dc_k", + .source_names = {"model.decoder.layers.cross_attn.k_proj"}, + .axes = {0, 1, 2}, + .shape = {layer_config.kv_heads, layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "dc_v", + .source_names = {"model.decoder.layers.cross_attn.v_proj"}, + .axes = {0, 1, 2}, + .shape = {layer_config.kv_heads, layer_config.qkv_dim, + config.model_dim}, + }); + Add(suffix, { + .base_name = "dc_att", + .source_names = {"model.decoder.layers.cross_attn.o_proj"}, + .preshape = {layer_config.heads, layer_config.qkv_dim, + config.model_dim}, + .axes = {0, 2, 1}, + .shape = {layer_config.heads, config.model_dim, + layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "dc_att_w", + .shape = {config.model_dim, + layer_config.heads * layer_config.qkv_dim}, + }); + Add(suffix, { + .base_name = "d_gate", + .source_names = {"model.decoder.layers.mlp.gating_einsum"}, + .axes = {0, 1, 2}, + .shape = {2, layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "d_gate1", + .shape = {layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "d_gate2", + .shape = {layer_config.ff_hidden_dim, config.model_dim}, + }); + Add(suffix, { + .base_name = "d_lin", + .source_names = {"model.decoder.layers.mlp.down_proj.weight"}, + .axes = {0, 1}, + .shape = {config.model_dim, layer_config.ff_hidden_dim}, + }); + Add(suffix, { + .base_name = "d_pre_sa", + .source_names = { + "model.decoder.layers.pre_self_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "d_post_sa", + .source_names = { + "model.decoder.layers.post_self_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "d_pre_ca", + .source_names = { + "model.decoder.layers.pre_cross_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "d_post_ca", + .source_names = { + "model.decoder.layers.post_cross_attn_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "d_pre_ff", + .source_names = { + "model.decoder.layers.pre_feedforward_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); + Add(suffix, { + .base_name = "d_post_ff", + .source_names = { + "model.decoder.layers.post_feedforward_layernorm.weight"}, + .axes = {0}, + .shape = {config.model_dim}, + .min_size = Type::kBF16, + }); +} + void TensorInfoRegistry::AddLayerTensors(const ModelConfig& config, const LayerConfig& layer_config, const size_t layer_idx) { @@ -658,10 +919,23 @@ TensorInfoRegistry::TensorInfoRegistry(const ModelConfig& config) { // in case those are changed without updating this. Better to allocate a bit // more than to 1.5-2x the size if too little. tensors_.reserve(10 + 32 * config.layer_configs.size() + - 24 * config.vit_config.layer_configs.size()); + 24 * config.vit_config.layer_configs.size() + + 8 * config.encoder_layer_configs.size() + + 14 * config.decoder_layer_configs.size()); AddModelTensors(config); - for (size_t i = 0; i < config.layer_configs.size(); ++i) { - AddLayerTensors(config, config.layer_configs[i], i); + if (config.is_encoder_decoder) { + AddT5GemmaModelTensors(config); + for (size_t i = 0; i < config.encoder_layer_configs.size(); ++i) { + AddT5GemmaEncoderLayerTensors(config, config.encoder_layer_configs[i], i); + } + for (size_t i = 0; i < config.decoder_layer_configs.size(); ++i) { + AddT5GemmaDecoderLayerTensors(config, config.decoder_layer_configs[i], i); + } + } + if (!config.is_encoder_decoder) { + for (size_t i = 0; i < config.layer_configs.size(); ++i) { + AddLayerTensors(config, config.layer_configs[i], i); + } } for (size_t i = 0; i < config.vit_config.layer_configs.size(); ++i) { AddImageLayerTensors(config, config.vit_config.layer_configs[i], i); diff --git a/gemma/tensor_info.h b/gemma/tensor_info.h index 97b88837..5d094eac 100644 --- a/gemma/tensor_info.h +++ b/gemma/tensor_info.h @@ -143,6 +143,13 @@ class TensorInfoRegistry { void AddModelTensors(const ModelConfig& config); void AddLayerTensors(const ModelConfig& config, const LayerConfig& layer_config, size_t layer_idx); + void AddT5GemmaModelTensors(const ModelConfig& config); + void AddT5GemmaEncoderLayerTensors(const ModelConfig& config, + const LayerConfig& layer_config, + size_t layer_idx); + void AddT5GemmaDecoderLayerTensors(const ModelConfig& config, + const LayerConfig& layer_config, + size_t layer_idx); void AddImageLayerTensors(const ModelConfig& config, const LayerConfig& layer_config, diff --git a/gemma/weights.cc b/gemma/weights.cc index 4bef99cb..72a1d000 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -164,6 +164,92 @@ void LayerWeightsPtrs::SplitAttW1() { qkv_einsum_w.SetPtr(nullptr, qkv_einsum_w.Cols()); } +static void InitAttWeightsGeneric(const LayerConfig& layer_config, + MatPtr& attn_vec_einsum_w, + MatPtr& att_weights, + std::vector& mat_owners, + const Allocator& allocator) { + HWY_ASSERT(attn_vec_einsum_w.HasPtr() ^ att_weights.HasPtr()); + if (att_weights.HasPtr() && !attn_vec_einsum_w.HasPtr()) return; + HWY_ASSERT(attn_vec_einsum_w.GetType() != Type::kNUQ); + + const size_t model_dim = layer_config.model_dim; + const size_t heads = layer_config.heads; + const size_t qkv_dim = layer_config.qkv_dim; + + att_weights.SetType(attn_vec_einsum_w.GetType()); + HWY_ASSERT(att_weights.Rows() == model_dim); + HWY_ASSERT(att_weights.Cols() == heads * qkv_dim); + HWY_ASSERT(attn_vec_einsum_w.Rows() == heads * model_dim); + HWY_ASSERT(attn_vec_einsum_w.Cols() == qkv_dim); + + { + static std::mutex m; + std::lock_guard lock(m); + mat_owners.push_back(MatOwner()); + mat_owners.back().AllocateFor(att_weights, allocator, MatPadding::kOdd); + } + + const size_t T_bytes = att_weights.ElementBytes(); + for (size_t m = 0; m < model_dim; ++m) { + uint8_t* HWY_RESTRICT out_row = att_weights.RowBytes(m); + for (size_t h = 0; h < heads; ++h) { + hwy::CopyBytes(attn_vec_einsum_w.RowBytes(h * model_dim + m), + out_row + h * qkv_dim * T_bytes, qkv_dim * T_bytes); + } + } + att_weights.SetScale(attn_vec_einsum_w.Scale()); +} + +static void SplitGateGeneric(const LayerConfig& layer_config, + MatPtr& gating_einsum_w, + MatPtr& gating_einsum_w1, + MatPtr& gating_einsum_w2) { + HWY_ASSERT(gating_einsum_w1.HasPtr() == gating_einsum_w2.HasPtr()); + HWY_ASSERT(gating_einsum_w.HasPtr() ^ gating_einsum_w1.HasPtr()); + if (gating_einsum_w1.HasPtr() && !gating_einsum_w.HasPtr()) return; + + const size_t ff_hidden_dim = layer_config.ff_hidden_dim; + HWY_ASSERT(gating_einsum_w.Rows() == 2 * ff_hidden_dim); + HWY_ASSERT(gating_einsum_w1.Rows() == ff_hidden_dim); + HWY_ASSERT(gating_einsum_w2.Rows() == ff_hidden_dim); + HWY_ASSERT(gating_einsum_w1.Cols() == gating_einsum_w.Cols()); + HWY_ASSERT(gating_einsum_w2.Cols() == gating_einsum_w.Cols()); + + const size_t stride = gating_einsum_w.Stride(); + gating_einsum_w1.SetPtr(gating_einsum_w.RowBytes(0), stride); + gating_einsum_w2.SetPtr(gating_einsum_w.RowBytes(ff_hidden_dim), stride); + gating_einsum_w1.SetType(gating_einsum_w.GetType()); + gating_einsum_w2.SetType(gating_einsum_w.GetType()); + gating_einsum_w1.SetScale(gating_einsum_w.Scale()); + gating_einsum_w2.SetScale(gating_einsum_w.Scale()); + gating_einsum_w.SetPtr(nullptr, gating_einsum_w.Cols()); +} + +static void SplitQKVGeneric(const LayerConfig& layer_config, + MatPtr& qkv_einsum_w, MatPtr& qkv_einsum_w1, + MatPtr& qkv_einsum_w2) { + HWY_ASSERT(qkv_einsum_w.HasPtr() ^ qkv_einsum_w1.HasPtr()); + if (qkv_einsum_w1.HasPtr() && !qkv_einsum_w.HasPtr()) return; + + const size_t w1_rows = layer_config.heads * layer_config.qkv_dim; + const size_t w2_rows = layer_config.kv_heads * 2 * layer_config.qkv_dim; + HWY_ASSERT(qkv_einsum_w.Rows() == w1_rows + w2_rows); + HWY_ASSERT(qkv_einsum_w1.Rows() == w1_rows); + HWY_ASSERT(qkv_einsum_w2.Rows() == w2_rows); + HWY_ASSERT(qkv_einsum_w1.Cols() == qkv_einsum_w.Cols()); + HWY_ASSERT(qkv_einsum_w2.Cols() == qkv_einsum_w.Cols()); + + const size_t stride = qkv_einsum_w.Stride(); + qkv_einsum_w1.SetPtr(qkv_einsum_w.RowBytes(0), stride); + qkv_einsum_w2.SetPtr(qkv_einsum_w.RowBytes(w1_rows), stride); + qkv_einsum_w1.SetType(qkv_einsum_w.GetType()); + qkv_einsum_w2.SetType(qkv_einsum_w.GetType()); + qkv_einsum_w1.SetScale(qkv_einsum_w.Scale()); + qkv_einsum_w2.SetScale(qkv_einsum_w.Scale()); + qkv_einsum_w.SetPtr(nullptr, qkv_einsum_w.Cols()); +} + static void HWY_MAYBE_UNUSED InitAttWeightsI8( const LayerConfig& layer_config, MatPtrT& attn_vec_einsum_w, MatPtrT& att_weights, std::vector& mat_owners, @@ -449,6 +535,27 @@ void LayerWeightsPtrs::Fixup(Model model, } } +void T5GemmaEncoderLayerWeightsPtrs::Fixup( + std::vector& mat_owners, ThreadingContext& ctx) { + InitAttWeightsGeneric(layer_config, attn_vec_einsum_w, att_weights, + mat_owners, ctx.allocator); + SplitGateGeneric(layer_config, gating_einsum_w, gating_einsum_w1, + gating_einsum_w2); + SplitQKVGeneric(layer_config, qkv_einsum_w, qkv_einsum_w1, qkv_einsum_w2); +} + +void T5GemmaDecoderLayerWeightsPtrs::Fixup( + std::vector& mat_owners, ThreadingContext& ctx) { + InitAttWeightsGeneric(layer_config, self_attn_vec_einsum_w, self_att_weights, + mat_owners, ctx.allocator); + InitAttWeightsGeneric(layer_config, cross_attn_vec_einsum_w, + cross_att_weights, mat_owners, ctx.allocator); + SplitGateGeneric(layer_config, gating_einsum_w, gating_einsum_w1, + gating_einsum_w2); + SplitQKVGeneric(layer_config, self_qkv_einsum_w, self_qkv_einsum_w1, + self_qkv_einsum_w2); +} + static void HWY_MAYBE_UNUSED InitAttWeightsNUQ( const LayerConfig& layer_config, MatPtrT& attn_vec_einsum_w, MatPtrT& att_weights, std::vector& mat_owners, @@ -518,6 +625,21 @@ void WeightsPtrs::CopyFrom(const WeightsPtrs& other) { void WeightsPtrs::Fixup(std::vector& mat_owners, ThreadingContext& ctx) { const size_t cluster_idx = 0; + if (config_.is_encoder_decoder) { + ParallelFor( + Parallelism::kFlat, t5gemma_encoder_layers.size(), ctx, cluster_idx, + Callers::kFixupWeights, [&](uint64_t layer, size_t /*worker*/) { + t5gemma_encoder_layers[layer].Fixup(mat_owners, ctx); + }); + + ParallelFor( + Parallelism::kFlat, t5gemma_decoder_layers.size(), ctx, cluster_idx, + Callers::kFixupWeights, [&](uint64_t layer, size_t /*worker*/) { + t5gemma_decoder_layers[layer].Fixup(mat_owners, ctx); + }); + return; + } + ParallelFor(Parallelism::kFlat, c_layers.size(), ctx, cluster_idx, Callers::kFixupWeights, [&](uint64_t layer, size_t /*worker*/) { GetLayer(layer)->Fixup(config_.model, mat_owners, ctx); diff --git a/gemma/weights.h b/gemma/weights.h index 20932d53..7f1fc0d8 100644 --- a/gemma/weights.h +++ b/gemma/weights.h @@ -328,6 +328,147 @@ struct LayerWeightsPtrs { void SplitAttW1(); }; +struct T5GemmaEncoderLayerWeightsPtrs { + T5GemmaEncoderLayerWeightsPtrs(size_t layer_idx, const LayerConfig& config, + const TensorInfoRegistry& tensors) + : layer_idx(layer_idx), + finder_(LayerSuffix(layer_idx), tensors), + qkv_einsum_w(finder_("e_qkv")), + qkv_einsum_w1(finder_("e_qkv1")), + qkv_einsum_w2(finder_("e_qkv2")), + attn_vec_einsum_w(finder_("e_att")), + att_weights(finder_("e_att_w")), + gating_einsum_w(finder_("e_gate")), + gating_einsum_w1(finder_("e_gate1")), + gating_einsum_w2(finder_("e_gate2")), + linear_w(finder_("e_lin")), + pre_attention_norm_scale(finder_("e_pre_att")), + post_attention_norm_scale(finder_("e_post_att")), + pre_ffw_norm_scale(finder_("e_pre_ff")), + post_ffw_norm_scale(finder_("e_post_ff")), + layer_config(config) {} + + const size_t layer_idx; + const MatFinder finder_; + + MatPtr qkv_einsum_w; + MatPtr qkv_einsum_w1; + MatPtr qkv_einsum_w2; + MatPtr attn_vec_einsum_w; + MatPtr att_weights; + MatPtr gating_einsum_w; + MatPtr gating_einsum_w1; + MatPtr gating_einsum_w2; + MatPtr linear_w; + MatPtr pre_attention_norm_scale; // at least BF16. + MatPtr post_attention_norm_scale; // at least BF16. + MatPtr pre_ffw_norm_scale; // at least BF16. + MatPtr post_ffw_norm_scale; // at least BF16. + + const LayerConfig& layer_config; + + template + void ForEachTensor(T5GemmaEncoderLayerWeightsPtrs* other1, + T5GemmaEncoderLayerWeightsPtrs* other2, Func func) { + func(TENSOR_ARGS(qkv_einsum_w, kMustRead)); + func(TENSOR_ARGS(qkv_einsum_w1, kMaybeRead)); + func(TENSOR_ARGS(qkv_einsum_w2, kMaybeRead)); + func(TENSOR_ARGS(attn_vec_einsum_w, kMustRead)); + func(TENSOR_ARGS(att_weights, kMaybeRead)); + func(TENSOR_ARGS(gating_einsum_w, kMustRead)); + func(TENSOR_ARGS(gating_einsum_w1, kMaybeRead)); + func(TENSOR_ARGS(gating_einsum_w2, kMaybeRead)); + func(TENSOR_ARGS(linear_w, kMustRead)); + func(TENSOR_ARGS(pre_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(post_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(pre_ffw_norm_scale, kMustRead)); + func(TENSOR_ARGS(post_ffw_norm_scale, kMustRead)); + } + + void Fixup(std::vector& mat_owners, ThreadingContext& ctx); +}; + +struct T5GemmaDecoderLayerWeightsPtrs { + T5GemmaDecoderLayerWeightsPtrs(size_t layer_idx, const LayerConfig& config, + const TensorInfoRegistry& tensors) + : layer_idx(layer_idx), + finder_(LayerSuffix(layer_idx), tensors), + self_qkv_einsum_w(finder_("d_qkv")), + self_qkv_einsum_w1(finder_("d_qkv1")), + self_qkv_einsum_w2(finder_("d_qkv2")), + self_attn_vec_einsum_w(finder_("d_att")), + self_att_weights(finder_("d_att_w")), + cross_q_einsum_w(finder_("dc_q")), + cross_k_einsum_w(finder_("dc_k")), + cross_v_einsum_w(finder_("dc_v")), + cross_attn_vec_einsum_w(finder_("dc_att")), + cross_att_weights(finder_("dc_att_w")), + gating_einsum_w(finder_("d_gate")), + gating_einsum_w1(finder_("d_gate1")), + gating_einsum_w2(finder_("d_gate2")), + linear_w(finder_("d_lin")), + pre_self_attention_norm_scale(finder_("d_pre_sa")), + post_self_attention_norm_scale(finder_("d_post_sa")), + pre_cross_attention_norm_scale(finder_("d_pre_ca")), + post_cross_attention_norm_scale(finder_("d_post_ca")), + pre_ffw_norm_scale(finder_("d_pre_ff")), + post_ffw_norm_scale(finder_("d_post_ff")), + layer_config(config) {} + + const size_t layer_idx; + const MatFinder finder_; + + MatPtr self_qkv_einsum_w; + MatPtr self_qkv_einsum_w1; + MatPtr self_qkv_einsum_w2; + MatPtr self_attn_vec_einsum_w; + MatPtr self_att_weights; + MatPtr cross_q_einsum_w; + MatPtr cross_k_einsum_w; + MatPtr cross_v_einsum_w; + MatPtr cross_attn_vec_einsum_w; + MatPtr cross_att_weights; + MatPtr gating_einsum_w; + MatPtr gating_einsum_w1; + MatPtr gating_einsum_w2; + MatPtr linear_w; + MatPtr pre_self_attention_norm_scale; // at least BF16. + MatPtr post_self_attention_norm_scale; // at least BF16. + MatPtr pre_cross_attention_norm_scale; // at least BF16. + MatPtr post_cross_attention_norm_scale; // at least BF16. + MatPtr pre_ffw_norm_scale; // at least BF16. + MatPtr post_ffw_norm_scale; // at least BF16. + + const LayerConfig& layer_config; + + template + void ForEachTensor(T5GemmaDecoderLayerWeightsPtrs* other1, + T5GemmaDecoderLayerWeightsPtrs* other2, Func func) { + func(TENSOR_ARGS(self_qkv_einsum_w, kMustRead)); + func(TENSOR_ARGS(self_qkv_einsum_w1, kMaybeRead)); + func(TENSOR_ARGS(self_qkv_einsum_w2, kMaybeRead)); + func(TENSOR_ARGS(self_attn_vec_einsum_w, kMustRead)); + func(TENSOR_ARGS(self_att_weights, kMaybeRead)); + func(TENSOR_ARGS(cross_q_einsum_w, kMustRead)); + func(TENSOR_ARGS(cross_k_einsum_w, kMustRead)); + func(TENSOR_ARGS(cross_v_einsum_w, kMustRead)); + func(TENSOR_ARGS(cross_attn_vec_einsum_w, kMustRead)); + func(TENSOR_ARGS(cross_att_weights, kMaybeRead)); + func(TENSOR_ARGS(gating_einsum_w, kMustRead)); + func(TENSOR_ARGS(gating_einsum_w1, kMaybeRead)); + func(TENSOR_ARGS(gating_einsum_w2, kMaybeRead)); + func(TENSOR_ARGS(linear_w, kMustRead)); + func(TENSOR_ARGS(pre_self_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(post_self_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(pre_cross_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(post_cross_attention_norm_scale, kMustRead)); + func(TENSOR_ARGS(pre_ffw_norm_scale, kMustRead)); + func(TENSOR_ARGS(post_ffw_norm_scale, kMustRead)); + } + + void Fixup(std::vector& mat_owners, ThreadingContext& ctx); +}; + // Holds layer-independent weight metadata and pointers plus per-layer // `LayerWeightsPtrs`. The tensor data is owned by `MatOwner`. struct WeightsPtrs { @@ -337,6 +478,10 @@ struct WeightsPtrs { finder_("", tensors_), // no suffix because these are per-model. embedder_input_embedding(finder_("c_embedding")), final_norm_scale(finder_("c_final_norm")), + t5gemma_encoder_embedding(finder_("enc_embedding")), + t5gemma_decoder_embedding(finder_("dec_embedding")), + t5gemma_encoder_final_norm_scale(finder_("enc_final_norm")), + t5gemma_decoder_final_norm_scale(finder_("dec_final_norm")), vit_encoder_norm_bias(finder_("enc_norm_bias")), vit_encoder_norm_scale(finder_("enc_norm_scale")), vit_img_embedding_bias(finder_("img_emb_bias")), @@ -346,10 +491,23 @@ struct WeightsPtrs { vit_img_head_kernel(finder_("img_head_kernel")), mm_embed_norm(finder_("mm_embed_norm")), c_layers() { - c_layers.reserve(config_.layer_configs.size()); - for (size_t idx = 0; idx < config_.layer_configs.size(); ++idx) { - const LayerConfig& layer_config = config_.layer_configs[idx]; - c_layers.emplace_back(idx, layer_config, tensors_); + if (config_.is_encoder_decoder) { + t5gemma_encoder_layers.reserve(config_.encoder_layer_configs.size()); + for (size_t idx = 0; idx < config_.encoder_layer_configs.size(); ++idx) { + const LayerConfig& layer_config = config_.encoder_layer_configs[idx]; + t5gemma_encoder_layers.emplace_back(idx, layer_config, tensors_); + } + t5gemma_decoder_layers.reserve(config_.decoder_layer_configs.size()); + for (size_t idx = 0; idx < config_.decoder_layer_configs.size(); ++idx) { + const LayerConfig& layer_config = config_.decoder_layer_configs[idx]; + t5gemma_decoder_layers.emplace_back(idx, layer_config, tensors_); + } + } else { + c_layers.reserve(config_.layer_configs.size()); + for (size_t idx = 0; idx < config_.layer_configs.size(); ++idx) { + const LayerConfig& layer_config = config_.layer_configs[idx]; + c_layers.emplace_back(idx, layer_config, tensors_); + } } for (size_t idx = 0; idx < config_.vit_config.layer_configs.size(); ++idx) { const LayerConfig& layer_config = config_.vit_config.layer_configs[idx]; @@ -368,6 +526,12 @@ struct WeightsPtrs { MatPtr embedder_input_embedding; MatPtr final_norm_scale; // at least BF16. + // T5Gemma text encoder-decoder parts. + MatPtr t5gemma_encoder_embedding; // at least BF16. + MatPtr t5gemma_decoder_embedding; // at least BF16. + MatPtr t5gemma_encoder_final_norm_scale; // at least BF16. + MatPtr t5gemma_decoder_final_norm_scale; // at least BF16. + // Vit parts. MatPtr vit_encoder_norm_bias; // at least BF16. MatPtr vit_encoder_norm_scale; // at least BF16. @@ -383,6 +547,8 @@ struct WeightsPtrs { std::vector c_layers; std::vector vit_layers; + std::vector t5gemma_encoder_layers; + std::vector t5gemma_decoder_layers; const LayerWeightsPtrs* GetLayer(size_t layer) const { return &c_layers[layer]; @@ -400,6 +566,34 @@ struct WeightsPtrs { void ForEachTensor(WeightsPtrs* other1, WeightsPtrs* other2, Func func) { LayerWeightsPtrs* other_layer1 = nullptr; LayerWeightsPtrs* other_layer2 = nullptr; + if (config_.is_encoder_decoder) { + func(TENSOR_ARGS(t5gemma_encoder_embedding, kMustRead)); + func(TENSOR_ARGS(t5gemma_decoder_embedding, kMustRead)); + func(TENSOR_ARGS(t5gemma_encoder_final_norm_scale, kMustRead)); + func(TENSOR_ARGS(t5gemma_decoder_final_norm_scale, kMustRead)); + + for (size_t layer_idx = 0; layer_idx < t5gemma_encoder_layers.size(); + ++layer_idx) { + auto* other_t5_layer1 = + other1 ? &other1->t5gemma_encoder_layers[layer_idx] : nullptr; + auto* other_t5_layer2 = + other2 ? &other2->t5gemma_encoder_layers[layer_idx] : nullptr; + t5gemma_encoder_layers[layer_idx].ForEachTensor( + other_t5_layer1, other_t5_layer2, func); + } + + for (size_t layer_idx = 0; layer_idx < t5gemma_decoder_layers.size(); + ++layer_idx) { + auto* other_t5_layer1 = + other1 ? &other1->t5gemma_decoder_layers[layer_idx] : nullptr; + auto* other_t5_layer2 = + other2 ? &other2->t5gemma_decoder_layers[layer_idx] : nullptr; + t5gemma_decoder_layers[layer_idx].ForEachTensor( + other_t5_layer1, other_t5_layer2, func); + } + return; + } + func(TENSOR_ARGS(embedder_input_embedding, kMustRead)); func(TENSOR_ARGS(final_norm_scale, kMustRead)); diff --git a/python/BUILD.bazel b/python/BUILD.bazel index 1c83bf67..3ff733f2 100644 --- a/python/BUILD.bazel +++ b/python/BUILD.bazel @@ -34,7 +34,6 @@ pybind_extension( py_binary( name = "run_example", srcs = ["run_example.py"], - strict_deps = False, deps = [ ":gemma", "@python_deps//absl_py", diff --git a/python/configs.cc b/python/configs.cc index 595ba69b..a3a21dd8 100644 --- a/python/configs.cc +++ b/python/configs.cc @@ -70,15 +70,17 @@ PYBIND11_MODULE(configs, py_module) { enum_(py_module, "PostQKType") .value("Rope", PostQKType::Rope) - .value("HalfRope", PostQKType::HalfRope); + .value("HalfRope", PostQKType::HalfRope) + .value("NormLocalRope", PostQKType::NormLocalRope); enum_(py_module, "ActivationType") .value("Gelu", ActivationType::Gelu); enum_(py_module, "QueryScaleType") .value("SqrtKeySize", QueryScaleType::SqrtKeySize) - .value("SqrtModelDimDivNumHeads", - QueryScaleType::SqrtModelDimDivNumHeads); + .value("SqrtModelDimDivNumHeads", + QueryScaleType::SqrtModelDimDivNumHeads) + .value("One", QueryScaleType::One); enum_(py_module, "ResidualType") .value("Add", ResidualType::Add); @@ -101,8 +103,10 @@ PYBIND11_MODULE(configs, py_module) { .value("GEMMA3_4B_LM", Model::GEMMA3_4B_LM) .value("GEMMA3_12B_LM", Model::GEMMA3_12B_LM) .value("GEMMA3_27B_LM", Model::GEMMA3_27B_LM) + .value("GEMMA4_26B_MOE", Model::GEMMA4_26B_MOE) + .value("T5GEMMA_S_S", Model::T5GEMMA_S_S) // Insert new models above this line. - .value("PALIGEMMA_448", Model::PALIGEMMA_448); + ; class_(py_module, "TensorInfo") .def(init()) @@ -142,6 +146,10 @@ PYBIND11_MODULE(configs, py_module) { .def_readwrite("activation", &LayerConfig::activation) .def_readwrite("post_qk", &LayerConfig::post_qk) .def_readwrite("use_qk_norm", &LayerConfig::use_qk_norm) + .def_readwrite("norm_v", &LayerConfig::norm_v) + .def_readwrite("num_experts", &LayerConfig::num_experts) + .def_readwrite("num_experts_per_datapoint", + &LayerConfig::num_experts_per_datapoint) .def_readwrite("internal", &LayerConfig::internal); class_(py_module, "VitConfig") @@ -185,6 +193,19 @@ PYBIND11_MODULE(configs, py_module) { .def_readwrite("internal", &ModelConfig::internal) .def_readwrite("use_global_timescale", &ModelConfig::use_global_timescale) + .def_readwrite("partial_rotary_factor", + &ModelConfig::partial_rotary_factor) + .def_readwrite("is_encoder_decoder", &ModelConfig::is_encoder_decoder) + .def_readwrite("encoder_num_layers", &ModelConfig::encoder_num_layers) + .def_readwrite("encoder_layer_configs", + &ModelConfig::encoder_layer_configs) + .def_readwrite("encoder_attention_window_sizes", + &ModelConfig::encoder_attention_window_sizes) + .def_readwrite("decoder_num_layers", &ModelConfig::decoder_num_layers) + .def_readwrite("decoder_layer_configs", + &ModelConfig::decoder_layer_configs) + .def_readwrite("decoder_attention_window_sizes", + &ModelConfig::decoder_attention_window_sizes) .def("add_layer_config", &ModelConfig::AddLayerConfig, arg("layer_config")) diff --git a/python/convert_from_safetensors.py b/python/convert_from_safetensors.py index 89e5e629..63ff06b2 100644 --- a/python/convert_from_safetensors.py +++ b/python/convert_from_safetensors.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Convert a PaliGemma[1/2] model from SafeTensors to gemma.cpp format.""" +"""Convert selected safetensors checkpoints to gemma.cpp SBS format. + +Supported families currently include PaliGemma, Gemma 3 LM variants, and +T5Gemma S/S. +""" # Tested with: # - PG1: huggingface.co/google/paligemma-3b-pt-224 # - PG1: huggingface.co/merve/paligemma_vqav2 @@ -29,6 +33,7 @@ from collections.abc import Sequence import csv +import glob import json import os import sys @@ -41,6 +46,19 @@ import safetensors import torch + +def _add_bazel_python_paths() -> None: + """Adds local Bazel extension outputs for direct script execution.""" + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + bazel_bins = [os.path.join(repo_root, "bazel-bin")] + bazel_bins += glob.glob(os.path.join(repo_root, "bazel-out", "*", "bin")) + for path in [repo_root] + bazel_bins: + if os.path.isdir(path) and path not in sys.path: + sys.path.insert(0, path) + + +_add_bazel_python_paths() + from compression.python import compression from python import configs @@ -70,6 +88,16 @@ def compute_scale(x: np.ndarray) -> float: return max(1.0, magnitude / 1.875) +def _get_safetensor_files(load_path: str) -> list[str]: + """Returns all safetensor files referenced by load_path.""" + if load_path.endswith(".json"): + with open(load_path, "r") as f: + j_obj = json.load(f) + files = list(set(j_obj["weight_map"].values())) + return [os.path.join(os.path.dirname(load_path), f) for f in files] + return [load_path] + + def _is_float_param(param_name: str) -> bool: """Returns whether the tensor should be stored as float32.""" for prefix in [ @@ -742,6 +770,307 @@ def add_gating_einsum(i): csv.writer(csv_handle).writerows(metadata) +def export_t5gemma_sbs( + model_specifier: str, + load_path: str, + tokenizer_file: str, + csv_file: str, + sbs_file: str, + weight_type: str = "bf16", +) -> None: + """Exports an sbs file from a T5Gemma safetensors checkpoint.""" + if not model_specifier.startswith("t5gemma-s-s"): + raise ValueError( + "Only T5Gemma S/S is currently supported by this converter path; got " + f"{model_specifier!r}." + ) + if weight_type not in ("bf16", "sfp"): + raise ValueError( + "T5Gemma weight_type must be 'bf16' or 'sfp'; got " + f"{weight_type!r}." + ) + requested_bf16 = "-bf16" in model_specifier + requested_sfp = "-sfp" in model_specifier + if requested_bf16 and weight_type != "bf16": + raise ValueError( + f"model_specifier {model_specifier!r} requests BF16, but " + f"--t5gemma_weight_type={weight_type} was passed." + ) + if requested_sfp and weight_type != "sfp": + raise ValueError( + f"model_specifier {model_specifier!r} requests SFP, but " + f"--t5gemma_weight_type={weight_type} was passed." + ) + + params: Dict[str, Any] = {} + for file in _get_safetensor_files(load_path): + with safetensors.safe_open(file, framework="pt") as f: + for k in f.keys(): + params[k] = f.get_tensor(k) + + config_path = os.path.join(os.path.dirname(load_path), "config.json") + hf_config = None + if os.path.exists(config_path): + with open(config_path, "r") as f: + hf_config = json.load(f) + if not hf_config.get("is_encoder_decoder", False): + raise ValueError("T5Gemma checkpoint config is not encoder-decoder.") + + enc_embed = params["model.encoder.embed_tokens.weight"] + dec_embed = params["model.decoder.embed_tokens.weight"] + vocab_size, model_dim = enc_embed.shape + assert dec_embed.shape == (vocab_size, model_dim) + + hidden_dim = params["model.encoder.layers.0.mlp.gate_proj.weight"].shape[0] + head_dim = ( + hf_config["encoder"]["head_dim"] + if hf_config is not None and "encoder" in hf_config + else 64 + ) + num_heads = ( + params["model.encoder.layers.0.self_attn.q_proj.weight"].shape[0] + // head_dim + ) + num_kv_heads = ( + params["model.encoder.layers.0.self_attn.k_proj.weight"].shape[0] + // head_dim + ) + num_encoder_layers = len( + set([ + k.split(".")[3] + for k in params + if k.startswith("model.encoder.layers.") + ]) + ) + num_decoder_layers = len( + set([ + k.split(".")[3] + for k in params + if k.startswith("model.decoder.layers.") + ]) + ) + + print( + f"T5Gemma: vocab={vocab_size} dim={model_dim} hidden={hidden_dim} " + f"heads={num_heads} kv_heads={num_kv_heads} head_dim={head_dim} " + f"enc_layers={num_encoder_layers} dec_layers={num_decoder_layers}" + ) + if ( + vocab_size != 256000 + or model_dim != 512 + or hidden_dim != 1024 + or num_heads != 8 + or num_kv_heads != 8 + or head_dim != 64 + or num_encoder_layers != 8 + or num_decoder_layers != 8 + ): + raise ValueError( + "Only T5Gemma S/S is currently supported; checkpoint dimensions were " + f"vocab={vocab_size}, dim={model_dim}, hidden={hidden_dim}, " + f"heads={num_heads}, kv_heads={num_kv_heads}, head_dim={head_dim}, " + f"enc_layers={num_encoder_layers}, dec_layers={num_decoder_layers}." + ) + + writer = compression.SbsWriter(sbs_file) + metadata = [] + scales = {} + + def t5gemma_tensor_type(sbs_name): + """Returns the storage type for a T5Gemma tensor.""" + if weight_type == "bf16": + return configs.Type.kBF16 + if ( + "embedding" in sbs_name + or "norm" in sbs_name + or "_pre_" in sbs_name + or "_post_" in sbs_name + ): + return configs.Type.kBF16 + return configs.Type.kSFP + + def add_data(param_name, data, expected_shape, sbs_name, layer_index=None): + if not isinstance(expected_shape, tuple): + expected_shape = (expected_shape,) + print(f"Writing {param_name} with shape {data.shape} e:{expected_shape}") + assert data.shape == expected_shape, param_name + assert isinstance(data, torch.Tensor) + data = data.to(torch.float32).numpy() + data = np.array(data) + + if layer_index is not None: + param_name = param_name % layer_index + sbs_name = sbs_name + f"_{layer_index}" + + value = flatten_f32(data) + scale = compute_scale(value) + both_names = param_name + "::" + sbs_name + metadata.append((both_names, data.dtype, data.shape, scale)) + packed = t5gemma_tensor_type(sbs_name) + if packed == configs.Type.kSFP: + assert scale == 1.0, f"Scale for {both_names} is not 1.0" + scales[sbs_name] = scale + + info = configs.TensorInfo() + info.name = sbs_name + info.shape = data.shape + writer.insert(sbs_name, value, packed, info) + sys.stdout.flush() + + def add_qkv(prefix, sbs_name, i): + q = params.pop(f"{prefix}.{i}.self_attn.q_proj.weight") + k = params.pop(f"{prefix}.{i}.self_attn.k_proj.weight") + v = params.pop(f"{prefix}.{i}.self_attn.v_proj.weight") + q = q.reshape(num_heads, head_dim, model_dim) + k = k.reshape(num_kv_heads, head_dim, model_dim) + v = v.reshape(num_kv_heads, head_dim, model_dim) + stacked = torch.stack((k, v), dim=0) + transposed = stacked.transpose(0, 1) + kv = transposed.reshape(2 * num_kv_heads, head_dim, model_dim) + qkv = torch.cat([q, kv], dim=0) + add_data( + f"{prefix}.%d.self_attn.qkv_proj.weight", + qkv, + (num_heads + 2 * num_kv_heads, head_dim, model_dim), + sbs_name, + i, + ) + + def add_att(prefix, sbs_name, i, attn_name="self_attn"): + o = params.pop(f"{prefix}.{i}.{attn_name}.o_proj.weight") + o = o.reshape(model_dim, num_heads, head_dim).permute(1, 0, 2) + add_data( + f"{prefix}.%d.{attn_name}.o_proj.weight", + o, + (num_heads, model_dim, head_dim), + sbs_name, + i, + ) + + def add_gating(prefix, sbs_name, i): + gate = params.pop(f"{prefix}.{i}.mlp.gate_proj.weight") + up = params.pop(f"{prefix}.{i}.mlp.up_proj.weight") + assert gate.shape == up.shape == (hidden_dim, model_dim) + gating = torch.stack([gate, up], dim=0) + add_data( + f"{prefix}.%d.mlp.gating_einsum.weight", + gating, + (2, hidden_dim, model_dim), + sbs_name, + i, + ) + + def add_cross_projection(prefix, hf_name, sbs_name, i): + w = params.pop(f"{prefix}.{i}.cross_attn.{hf_name}_proj.weight") + heads = num_heads if hf_name == "q" else num_kv_heads + w = w.reshape(heads, head_dim, model_dim) + add_data( + f"{prefix}.%d.cross_attn.{hf_name}_proj.weight", + w, + (heads, head_dim, model_dim), + sbs_name, + i, + ) + + add_data( + "model.encoder.embed_tokens.weight", + params.pop("model.encoder.embed_tokens.weight"), + (vocab_size, model_dim), + "enc_embedding", + ) + add_data( + "model.decoder.embed_tokens.weight", + params.pop("model.decoder.embed_tokens.weight"), + (vocab_size, model_dim), + "dec_embedding", + ) + add_data( + "model.encoder.norm.weight", + params.pop("model.encoder.norm.weight"), + (model_dim,), + "enc_final_norm", + ) + add_data( + "model.decoder.norm.weight", + params.pop("model.decoder.norm.weight"), + (model_dim,), + "dec_final_norm", + ) + + enc_prefix = "model.encoder.layers" + for i in range(num_encoder_layers): + add_qkv(enc_prefix, "e_qkv", i) + add_att(enc_prefix, "e_att", i) + add_gating(enc_prefix, "e_gate", i) + add_data( + f"{enc_prefix}.%d.mlp.down_proj.weight", + params.pop(f"{enc_prefix}.{i}.mlp.down_proj.weight"), + (model_dim, hidden_dim), + "e_lin", + i, + ) + for hf_name, sbs_name in [ + ("pre_self_attn_layernorm", "e_pre_att"), + ("post_self_attn_layernorm", "e_post_att"), + ("pre_feedforward_layernorm", "e_pre_ff"), + ("post_feedforward_layernorm", "e_post_ff"), + ]: + add_data( + f"{enc_prefix}.%d.{hf_name}.weight", + params.pop(f"{enc_prefix}.{i}.{hf_name}.weight"), + (model_dim,), + sbs_name, + i, + ) + + dec_prefix = "model.decoder.layers" + for i in range(num_decoder_layers): + add_qkv(dec_prefix, "d_qkv", i) + add_att(dec_prefix, "d_att", i) + add_cross_projection(dec_prefix, "q", "dc_q", i) + add_cross_projection(dec_prefix, "k", "dc_k", i) + add_cross_projection(dec_prefix, "v", "dc_v", i) + add_att(dec_prefix, "dc_att", i, attn_name="cross_attn") + add_gating(dec_prefix, "d_gate", i) + add_data( + f"{dec_prefix}.%d.mlp.down_proj.weight", + params.pop(f"{dec_prefix}.{i}.mlp.down_proj.weight"), + (model_dim, hidden_dim), + "d_lin", + i, + ) + for hf_name, sbs_name in [ + ("pre_self_attn_layernorm", "d_pre_sa"), + ("post_self_attn_layernorm", "d_post_sa"), + ("pre_cross_attn_layernorm", "d_pre_ca"), + ("post_cross_attn_layernorm", "d_post_ca"), + ("pre_feedforward_layernorm", "d_pre_ff"), + ("post_feedforward_layernorm", "d_post_ff"), + ]: + add_data( + f"{dec_prefix}.%d.{hf_name}.weight", + params.pop(f"{dec_prefix}.{i}.{hf_name}.weight"), + (model_dim,), + sbs_name, + i, + ) + + assert not params, "Some params were not used: %s" % params.keys() + + sbs_model_specifier = model_specifier + if sbs_model_specifier == "t5gemma-s-s": + sbs_model_specifier = f"t5gemma-s-s-{weight_type}" + sbs_config = configs.ModelConfig(sbs_model_specifier) + if not sbs_config.is_encoder_decoder: + raise ValueError( + f"{sbs_model_specifier!r} is not an encoder-decoder config." + ) + writer.write(sbs_config, tokenizer_file) + + with open(csv_file, "w") as csv_handle: + csv.writer(csv_handle).writerows(metadata) + + _MODEL_SPECIFIER = flags.DEFINE_string( "model_specifier", None, @@ -770,6 +1099,13 @@ def add_gating_einsum(i): "/tmp/gemmacpp.sbs", "Path to the sbs file to write", ) +_T5GEMMA_WEIGHT_TYPE = flags.DEFINE_enum( + "t5gemma_weight_type", + "bf16", + ["bf16", "sfp"], + "For T5Gemma only, choose BF16 for parity/debug correctness or SFP for a " + "smaller experimental mixed BF16/SFP file.", +) def main(argv: Sequence[str]) -> None: @@ -782,6 +1118,7 @@ def main(argv: Sequence[str]) -> None: tokenizer_file = _TOKENIZER_FILE.value metadata_file = _METADATA_FILE.value sbs_file = _SBS_FILE.value + t5gemma_weight_type = _T5GEMMA_WEIGHT_TYPE.value logging.info( "\n====\nReading %s from %s and %s, writing to %s\n====", @@ -798,10 +1135,19 @@ def main(argv: Sequence[str]) -> None: export_gemma3_lm_sbs( model_specifier, load_path, tokenizer_file, metadata_file, sbs_file ) + elif model_specifier.startswith("t5gemma"): + export_t5gemma_sbs( + model_specifier, + load_path, + tokenizer_file, + metadata_file, + sbs_file, + weight_type=t5gemma_weight_type, + ) else: raise app.UsageError( f"Unsupported model_specifier {model_specifier!r}. Expected a " - "'paligemma*' or 'gemma3-*-lm-*' specifier." + "'paligemma*', 'gemma3-*-lm-*', or 't5gemma*' specifier." )