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/

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,10 @@ 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]),
+        num_layers(config.is_encoder_decoder ? config.decoder_num_layers
+                                             : config.num_layers),
 
         x(MatFactory("x", batch_size, config.model_dim, ctx.allocator)),
         x_bf(MatFactory("x_bf", batch_size, config.model_dim, ctx.allocator)),
@@ -366,28 +369,76 @@ struct Activations {
                       ctx.allocator)),
         ffw_out(
             MatFactory("ffw_out", batch_size, config.model_dim, ctx.allocator)),
+        t5_encoder_pre_att_rms_out(MatFactory(
+            "t5_e_pre_att", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.model_dim : 0, ctx.allocator)),
+        t5_encoder_q(MatFactory(
+            "t5_e_q", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.encoder_layer_configs[0].heads *
+                                            config.encoder_layer_configs[0].qkv_dim
+                                      : 0,
+            ctx.allocator)),
+        t5_encoder_kv(MatFactory(
+            "t5_e_kv", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder
+                ? 2 * config.encoder_layer_configs[0].kv_heads *
+                      config.encoder_layer_configs[0].qkv_dim
+                : 0,
+            ctx.allocator)),
+        t5_encoder_att_out(MatFactory(
+            "t5_e_att_out", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.encoder_layer_configs[0].heads *
+                                            config.encoder_layer_configs[0].qkv_dim
+                                      : 0,
+            ctx.allocator)),
+        t5_encoder_att_sums(MatFactory(
+            "t5_e_att_sums", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.model_dim : 0, ctx.allocator)),
+        t5_encoder_pre_ffw_rms_out(MatFactory(
+            "t5_e_pre_ffw", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.model_dim : 0, ctx.allocator)),
+        t5_encoder_c1(MatFactory(
+            "t5_e_c1", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.encoder_layer_configs[0].ff_hidden_dim
+                                      : 0,
+            ctx.allocator)),
+        t5_encoder_c2(MatFactory(
+            "t5_e_c2", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.encoder_layer_configs[0].ff_hidden_dim
+                                      : 0,
+            ctx.allocator)),
+        t5_encoder_ffw_out(MatFactory(
+            "t5_e_ffw_out", config.is_encoder_decoder ? seq_len : 0,
+            config.is_encoder_decoder ? config.model_dim : 0, ctx.allocator)),
+        t5_encoder_inv_timescale(
+            config.is_encoder_decoder
+                ? CreateInvTimescale(
+                      ctx.allocator, config.encoder_layer_configs[0].qkv_dim,
+                      config.encoder_layer_configs[0].post_qk ==
+                          PostQKType::HalfRope)
+                : MatStorageT<float>()),
 
         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(num_layers, max_workers),
+        s_ffw_hidden(num_layers, max_workers),
+        s_ffw_out(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(num_layers, max_workers),
+        s_router_logits(num_layers, max_workers),
+        s_expert_in(num_layers, max_workers),
+        s_expert_hidden(num_layers, max_workers),
+        s_expert_out(num_layers, max_workers),
+        s_w_expert_in1(num_layers, max_workers),
+        s_w_expert_in2(num_layers, max_workers),
+        s_w_expert_hidden(num_layers, max_workers),
+        s_w_gating_einsum_w1(num_layers, max_workers),
+        s_w_gating_einsum_w2(num_layers, max_workers),
+        s_w_linear_w(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,
@@ -405,6 +456,16 @@ struct Activations {
     C2.AllocateAndAttachRowPtrs(row_ptrs);
     ffw_out.AllocateAndAttachRowPtrs(row_ptrs);
 
+    if (config.is_encoder_decoder) {
+      t5_encoder_q.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_kv.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_att_out.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_att_sums.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_c1.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_c2.AllocateAndAttachRowPtrs(row_ptrs);
+      t5_encoder_ffw_out.AllocateAndAttachRowPtrs(row_ptrs);
+    }
+
     if (layer_config.IsMoE()) {
       router_logits.AllocateAndAttachRowPtrs(row_ptrs);
 
@@ -469,6 +530,7 @@ struct Activations {
   }
 
   const LayerConfig& layer_config;
+  const size_t num_layers;
 
   MatStorageT<float> x;    // input
   MatStorageT<BF16> x_bf;  // output of final RMSNorm, input to EmbeddingMatmul
@@ -481,6 +543,30 @@ struct Activations {
   MatStorageT<BF16> C2;
   MatStorageT<float> ffw_out;
 
+  void SetT5EncoderSourceLen(size_t source_len) {
+    t5_encoder_pre_att_rms_out.OverrideRows(source_len);
+    t5_encoder_q.OverrideRows(source_len);
+    t5_encoder_kv.OverrideRows(source_len);
+    t5_encoder_att_out.OverrideRows(source_len);
+    t5_encoder_att_sums.OverrideRows(source_len);
+    t5_encoder_pre_ffw_rms_out.OverrideRows(source_len);
+    t5_encoder_c1.OverrideRows(source_len);
+    t5_encoder_c2.OverrideRows(source_len);
+    t5_encoder_ffw_out.OverrideRows(source_len);
+  }
+
+  // T5Gemma encoder scratch storage.
+  MatStorageT<BF16> t5_encoder_pre_att_rms_out;
+  MatStorageT<float> t5_encoder_q;
+  MatStorageT<float> t5_encoder_kv;
+  MatStorageT<float> t5_encoder_att_out;
+  MatStorageT<float> t5_encoder_att_sums;
+  MatStorageT<BF16> t5_encoder_pre_ffw_rms_out;
+  MatStorageT<BF16> t5_encoder_c1;
+  MatStorageT<BF16> t5_encoder_c2;
+  MatStorageT<float> t5_encoder_ffw_out;
+  MatStorageT<float> t5_encoder_inv_timescale;
+
   const size_t max_workers;
   TensorStats s_ffw_in;
   TensorStats s_ffw_hidden;  // after Activation+gating

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;