Add initial T5Gemma S/S support

.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
+.env.local
+.venv
+
+# OS Files
+.DS_Store
+
+# Temporary
+.temp/

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",
     ],
 )

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.
 

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

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(

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 <typename Packed>
@@ -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 {

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<hwy::AlignedFreeUniquePtr<uint8_t*[]>>& 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,

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<int>(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<int>(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;
 

gemma/configs.h

@@ -136,7 +136,7 @@ enum class PostQKType {
 
 static inline bool EnumValid(PostQKType type) {
   return static_cast<size_t>(type) <
-         static_cast<size_t>(PostNormType::kSentinel);
+         static_cast<size_t>(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<LayerConfig> encoder_layer_configs;
+  std::vector<uint32_t> encoder_attention_window_sizes;
+  uint32_t decoder_num_layers = 0;
+  std::vector<LayerConfig> decoder_layer_configs;
+  std::vector<uint32_t> 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`.