[gemma_cpp] Add support for Qwen3 models (0.6B, 1.7B and 4B).

deepseek/deepseek_tensors.cc

@@ -58,16 +58,6 @@ void TensorInfoRegistry::AddDeepSeekModelTensors(const ModelConfig& config) {
                    });
   };
 
-  if (config.HasMLA()) {
-    // DeepSeek models have an untied output head.
-    Add(no_suffix, {
-                       .base_name = "lm_head",
-                       .source_names = {"lm_head.weight", "lm_head"},
-                       .axes = {0, 1},
-                       .shape = {config.vocab_size, config.model_dim},
-                       .min_size = Type::kBF16,
-                   });
-  }
   if (config.hc_mult > 1) {
     add_hc_collapse("hc_head", "");
   }

gemma/configs.cc

@@ -666,6 +666,85 @@ LayerConfig ModelConfig::MTPLayerConfig() const {
   return lc;
 }
 
+static ModelConfig ConfigBaseQwen3() {
+  ModelConfig config = ConfigNoSSM();
+  config.vocab_size = 151936;
+  config.max_seq_len = 32768;
+  config.eos_id = 151645;
+  config.secondary_eos_id = 151643;
+  return config;
+}
+
+static LayerConfig LayerConfigQwen3(size_t model_dim, size_t ff_hidden_dim,
+                                    size_t heads, size_t kv_heads,
+                                    size_t qkv_dim) {
+  LayerConfig config;
+  config.model_dim = model_dim;
+  config.ff_hidden_dim = ff_hidden_dim;
+  config.heads = heads;
+  config.kv_heads = kv_heads;
+  config.qkv_dim = qkv_dim;
+  config.optimized_gating = false;
+  config.post_norm = PostNormType::None;
+  config.activation = ActivationType::Silu;
+  config.use_qk_norm = true;
+  return config;
+}
+
+static ModelConfig ConfigQWEN3_600M() {
+  ModelConfig config = ConfigBaseQwen3();
+  config.display_name = "Qwen3_0.6B";
+  config.model = Model::QWEN3_600M;
+  config.wrapping = PromptWrapping::GEMMA_IT;
+  config.model_dim = 1024;
+
+  LayerConfig layer_config =
+      LayerConfigQwen3(config.model_dim, 3072, 16, 8, 128);
+  config.num_layers = 28;
+  config.layer_configs = {config.num_layers, layer_config};
+  config.query_scale = QueryScaleType::SqrtKeySize;
+  config.use_global_timescale = true;
+  config.attention_window_sizes =
+      FixedAttentionWindowSizes<28>(config.max_seq_len);
+  return config;
+}
+
+static ModelConfig ConfigQWEN3_2B() {
+  ModelConfig config = ConfigBaseQwen3();
+  config.display_name = "Qwen3_1.7B";
+  config.model = Model::QWEN3_2B;
+  config.wrapping = PromptWrapping::GEMMA_IT;
+  config.model_dim = 2048;
+
+  LayerConfig layer_config =
+      LayerConfigQwen3(config.model_dim, 6144, 16, 8, 128);
+  config.num_layers = 28;
+  config.layer_configs = {config.num_layers, layer_config};
+  config.query_scale = QueryScaleType::SqrtKeySize;
+  config.use_global_timescale = true;
+  config.attention_window_sizes =
+      FixedAttentionWindowSizes<28>(config.max_seq_len);
+  return config;
+}
+
+static ModelConfig ConfigQwen3_4B() {
+  ModelConfig config = ConfigBaseQwen3();
+  config.display_name = "Qwen3_4B";
+  config.model = Model::QWEN3_4B;
+  config.wrapping = PromptWrapping::GEMMA_IT;
+  config.model_dim = 2560;
+
+  LayerConfig layer_config =
+      LayerConfigQwen3(config.model_dim, 9728, 32, 8, 128);
+  config.num_layers = 36;
+  config.layer_configs = {config.num_layers, layer_config};
+  config.query_scale = QueryScaleType::SqrtKeySize;
+  config.use_global_timescale = true;
+  config.attention_window_sizes =
+      FixedAttentionWindowSizes<36>(config.max_seq_len);
+  return config;
+}
+
 static ModelConfig ConfigFromModel(Model model) {
   switch (model) {
     case Model::GEMMA2_2B:
@@ -704,6 +783,12 @@ static ModelConfig ConfigFromModel(Model model) {
       return ConfigGemma4_2B();
     case Model::DEEPSEEK4_FLASH:
       return ConfigDeepSeek4_Flash();
+    case Model::QWEN3_600M:
+      return ConfigQWEN3_600M();
+    case Model::QWEN3_2B:
+      return ConfigQWEN3_2B();
+    case Model::QWEN3_4B:
+      return ConfigQwen3_4B();
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -749,6 +834,12 @@ const char* ModelPrefix(Model model) {
       return "gemma4-2b";
     case Model::DEEPSEEK4_FLASH:
       return "deepseek4-flash";
+    case Model::QWEN3_600M:
+      return "qwen3-0_6b";
+    case Model::QWEN3_2B:
+      return "qwen3-2b";
+    case Model::QWEN3_4B:
+      return "qwen3-4b";
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -943,6 +1034,13 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
     case 27:
       return (layer_types & kDeduced448) ? Model::PALIGEMMA2_3B_448
                                          : Model::PALIGEMMA2_3B_224;
+    case 28:
+      if (blob_path.path.find("qwen3-2b") != std::string::npos ||
+          blob_path.path.find("qwen3-1_7b") != std::string::npos) {
+        return Model::QWEN3_2B;
+      }
+      return Model::QWEN3_600M;
+
     case 30:
       return Model::GEMMA4_26B_MOE;
 
@@ -951,6 +1049,8 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
                                          : Model::GEMMA3_4B_LM;
     case 35:
       return Model::GEMMA4_2B;
+    case 36:
+      return Model::QWEN3_4B;
     case 42:
       if (layer_types & kDeducedViT) {
         return (layer_types & kDeduced448) ? Model::PALIGEMMA2_10B_448

gemma/configs.h

@@ -273,6 +273,9 @@ enum class Model {
   GEMMA4_26B_MOE,
   GEMMA4_2B,
   DEEPSEEK4_FLASH,
+  QWEN3_600M,
+  QWEN3_2B,  // 1.7B rounded up for readability.
+  QWEN3_4B,
   kSentinel,
 };
 
@@ -632,6 +635,15 @@ struct ModelConfig : public IFields {
     return false;
   }
 
+  bool IsQwen3() const {
+    return model == Model::QWEN3_600M || model == Model::QWEN3_2B ||
+           model == Model::QWEN3_4B;
+  }
+
+  bool HasLmHead() const {
+    return model == Model::QWEN3_600M || model == Model::QWEN3_2B || HasMLA();
+  }
+
   // Synthesized config for the multi-token-prediction block (DeepSeek V4):
   // a full extra layer (dense MLA + MoE) used for speculative decoding, not
   // part of the main stack. Only valid when `num_mtp_layers > 0`.

gemma/gemma.cc

@@ -220,6 +220,10 @@ static float EmbeddingScaling(size_t model_dim) {
       hwy::ConvertScalarTo<BF16>(sqrtf(static_cast<float>(model_dim))));
 }
 
+static bool HasEmbeddingScaling(const ModelConfig& model_config) {
+  return !(model_config.IsQwen3() || model_config.HasMLA());
+}
+
 // `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
@@ -256,9 +260,9 @@ HWY_NOINLINE size_t EmbedMMToken(int token, size_t x_row, size_t pos,
   }
 
   const size_t model_dim = model_config.model_dim;
-  // DeepSeek does not scale embeddings by sqrt(model_dim).
+  // Qwen3/DeepSeekV4 does not scale embeddings by sqrt(model_dim).
   const float emb_scaling =
-      model_config.HasMLA() ? 1.0f : EmbeddingScaling(model_dim);
+      HasEmbeddingScaling(model_config) ? EmbeddingScaling(model_dim) : 1.0f;
 
   HWY_DASSERT(token >= 0);
   HWY_DASSERT(token < static_cast<int>(model_config.vocab_size));

gemma/tensor_info.cc

@@ -59,6 +59,16 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) {
                      .min_size = Type::kBF16,
                  });
   AddDeepSeekModelTensors(config);
+  if (config.HasLmHead()) {
+    Add(no_suffix,
+        {
+            .base_name = "lm_head",
+            .source_names = {"lm_head/weight", "lm_head.weight", "lm_head"},
+            .axes = {0, 1},
+            .shape = {config.vocab_size, config.model_dim},
+            .min_size = Type::kBF16,
+        });
+  }
   Add(no_suffix, {
                      .base_name = "enc_norm_bias",
                      .source_names = {"img/Transformer/encoder_norm/bias"},

gemma/tokenizer.cc

@@ -75,6 +75,12 @@ GemmaChatTemplate::GemmaChatTemplate(const GemmaTokenizer& tokenizer,
     sot_user_ = {105, 2364, 107};
     sot_model_ = {105, 4368, 107};
     eot_ = {106, 107};
+  } else if (model == Model::QWEN3_600M || model == Model::QWEN3_2B ||
+             model == Model::QWEN3_4B) {
+    prepend_bos_ = false;
+    HWY_ASSERT(tokenizer.Encode("<|im_start|>user\n", &sot_user_));
+    HWY_ASSERT(tokenizer.Encode("<|im_start|>assistant\n", &sot_model_));
+    HWY_ASSERT(tokenizer.Encode("<|im_end|>\n", &eot_));
   } else {
     sot_user_.reserve(3);
     if (!tokenizer.Encode("<start_of_turn>user\n", &sot_user_)) return;
@@ -101,7 +107,9 @@ std::vector<int> GemmaChatTemplate::Apply(size_t pos,
 
   // Start with BOS, or prepend end_of_turn if this is a continuation.
   if (pos == 0) {
-    out.push_back(BOS_ID);
+    if (prepend_bos_) {
+      out.push_back(BOS_ID);
+    }
   } else {
     out.insert(out.cend(), eot_.cbegin(), eot_.cend());
   }

gemma/tokenizer.h

@@ -77,6 +77,7 @@ class GemmaChatTemplate {
   std::vector<int> pali_sep_;
   std::vector<int> vlm_soi_;
   std::vector<int> vlm_eoi_;
+  bool prepend_bos_ = true;
 };
 
 std::vector<int> WrapAndTokenize(const GemmaTokenizer& tokenizer,

gemma/weights.h

@@ -16,7 +16,6 @@
 #ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_H_
 #define THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_H_
 
-#include <math.h>  // isnan
 #include <stddef.h>
 #include <stdint.h>
 
@@ -252,7 +251,7 @@ struct LayerWeightsPtrs {
 
   MatPtr key_norm_scale;    // at least BF16.
   MatPtr query_norm_scale;  // at least BF16.
-  
+
   MatPtr router_scale;
   MatPtr p_expert_sc;
   MatPtr post_ffw1_ns;
@@ -594,7 +593,7 @@ struct WeightsPtrs {
     LayerWeightsPtrs* other_layer2 = nullptr;
     func(TENSOR_ARGS(embedder_input_embedding, kMustRead));
     func(TENSOR_ARGS(final_norm_scale, kMustRead));
-    if (config_.HasMLA()) {
+    if (config_.HasLmHead()) {
       func(TENSOR_ARGS(lm_head, kMustRead));
     }
     if (config_.hc_mult > 1) {

python/configs.cc

@@ -104,6 +104,9 @@ PYBIND11_MODULE(configs, py_module) {
       .value("GEMMA3_12B_LM", gcpp::Model::GEMMA3_12B_LM)
       .value("GEMMA3_27B_LM", gcpp::Model::GEMMA3_27B_LM)
       .value("DEEPSEEK4_FLASH", gcpp::Model::DEEPSEEK4_FLASH)
+      .value("QWEN3_600M", gcpp::Model::QWEN3_600M)
+      .value("QWEN3_2B", gcpp::Model::QWEN3_2B)
+      .value("QWEN3_4B", gcpp::Model::QWEN3_4B)
       // Insert new models above this line.
   .value("PALIGEMMA_448", gcpp::Model::PALIGEMMA_448);