Add Gemma4 multimodal support (vision + audio)

cmake/deps.txt

@@ -14,7 +14,7 @@ pybind11;https://github.com/pybind/pybind11/archive/refs/tags/v2.13.6.zip;f78029
 googletest;https://github.com/google/googletest/archive/530d5c8c84abd2a46f38583ee817743c9b3a42b4.zip;5e3a61db2aa975cfd0f97ba92c818744e7fa7034
 microsoft_wil;https://github.com/microsoft/wil/archive/refs/tags/v1.0.230629.1.zip;e4a542a323c070376f7c2d1973d0f7ddbc1d2fa5
 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e
-onnxruntime_extensions;https://github.com/microsoft/onnxruntime-extensions.git;539d380ce9c2fcdfc9fd9f151ef5604425215aa9
+onnxruntime_extensions;https://github.com/microsoft/onnxruntime-extensions.git;e094cc816679d0b2b5fe2b4fd7f73e5b1844b425
 
 # These two dependencies are for the optional constrained decoding feature (USE_GUIDANCE)
 llguidance;https://github.com/microsoft/llguidance.git;94fa39128ef184ffeda33845f6d333f332a34b4d

examples/python/common.py

@@ -284,9 +284,10 @@ def get_user_content(model_type: str, num_images: int, num_audios: int, prompt:
         image_tags = "".join(["[IMG]" for _ in range(num_images)])
         content = image_tags + prompt
     else:
-        # Gemma-3 style: structured content
+        # Gemma-3/4 style: structured content with image and audio entries
         image_tags = [{"type": "image"} for _ in range(num_images)]
-        content = image_tags + [{"type": "text", "text": prompt}]
+        audio_tags = [{"type": "audio"} for _ in range(num_audios)]
+        content = image_tags + audio_tags + [{"type": "text", "text": prompt}]
     return content
 
 @dataclass

src/config.cpp

@@ -648,6 +648,8 @@ struct VisionInputs_Element : JSON::Element {
   void OnValue(std::string_view name, JSON::Value value) override {
     if (name == "pixel_values") {
       v_.pixel_values = JSON::Get<std::string_view>(value);
+    } else if (name == "pixel_position_ids") {
+      v_.pixel_position_ids = JSON::Get<std::string_view>(value);
     } else if (name == "image_sizes") {
       v_.image_sizes = JSON::Get<std::string_view>(value);
     } else if (name == "image_grid_thw") {
@@ -1096,6 +1098,10 @@ struct Model_Element : JSON::Element {
       v_.sep_token_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "image_token_id") {
       v_.image_token_id = static_cast<int>(JSON::Get<double>(value));
+    } else if (name == "audio_token_id") {
+      v_.audio_token_id = static_cast<int>(JSON::Get<double>(value));
+    } else if (name == "boa_token_id") {
+      v_.boa_token_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "video_token_id") {
       v_.video_token_id = static_cast<int>(JSON::Get<double>(value));
     } else if (name == "vision_start_token_id") {

src/config.h

@@ -40,6 +40,7 @@ struct Config {
     static constexpr std::string_view ImageSizesName = "image_sizes";
     static constexpr std::string_view ImageGridThwName = "image_grid_thw";
     static constexpr std::string_view ImageAttentionMaskName = "image_attention_mask";
+    static constexpr std::string_view PixelPositionIdsName = "pixel_position_ids";
     static constexpr std::string_view ImageFeaturesName = "image_features";
     static constexpr std::string_view NumImageTokens = "num_image_tokens";
 
@@ -129,6 +130,8 @@ struct Config {
 
     // Qwen2.5-VL specific token IDs
     int image_token_id{};
+    int audio_token_id{};
+    int boa_token_id{};
     int video_token_id{};
     int vision_start_token_id{};
 
@@ -237,6 +240,7 @@ struct Config {
 
       struct Inputs {
         std::string pixel_values{Defaults::PixelValuesName};
+        std::string pixel_position_ids{Defaults::PixelPositionIdsName};
         std::string image_sizes{Defaults::ImageSizesName};
         std::string image_grid_thw{Defaults::ImageSizesName};          // Qwen2.5-VL uses image_grid_thw, defaults to image_sizes
         std::string attention_mask{Defaults::ImageAttentionMaskName};  // image attention mask

src/models/gemma4_multimodal_processor.h

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#pragma once
+
+#include "processor.h"
+
+namespace Generators {
+
+struct Gemma4MultiModalProcessor : Processor {
+  Gemma4MultiModalProcessor(Config& config, const SessionInfo& session_info);
+
+  virtual std::unique_ptr<NamedTensors> Process(const Tokenizer& tokenizer, const Payload& payload) const override;
+
+ private:
+  ort_extensions::OrtxObjectPtr<OrtxProcessor> image_processor_;
+  ort_extensions::OrtxObjectPtr<OrtxFeatureExtractor> audio_processor_;
+
+  ONNXTensorElementDataType pixel_values_type_;
+  ONNXTensorElementDataType audio_features_type_;
+
+  bool has_speech_{false};
+  size_t vision_soft_tokens_per_image_{260};
+};
+
+}  // namespace Generators

src/models/kv_cache.cpp

@@ -207,6 +207,49 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state)
   type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]);
   empty_past_ = OrtValue::CreateTensor(Allocator(), shape_, type_);
 
+  // Auto-detect per-layer head_dim from ONNX session input shapes.
+  // Models like Gemma 4 have dual head_dim: sliding-window layers use head_dim=256,
+  // full-attention layers use global_head_dim=512.
+  {
+    bool has_varying_head_dim = false;
+    std::vector<int64_t> per_layer_head_dim(layer_count_, shape_[3]);
+    for (int i = 0; i < layer_count_; ++i) {
+      auto input_shape = model_.session_info_.GetInputShape(input_name_strings_[i * 2]);
+      if (!input_shape.empty()) {
+        int64_t layer_head_dim = input_shape.back();
+        if (layer_head_dim > 0 && layer_head_dim != shape_[3]) {
+          has_varying_head_dim = true;
+        }
+        if (layer_head_dim > 0) {
+          per_layer_head_dim[i] = layer_head_dim;
+        }
+      }
+    }
+    if (has_varying_head_dim) {
+      if (layer_shapes_.empty()) {
+        layer_shapes_.resize(layer_count_);
+        for (int i = 0; i < layer_count_; ++i) {
+          layer_shapes_[i] = shape_;
+        }
+      }
+      for (int i = 0; i < layer_count_; ++i) {
+        layer_shapes_[i][3] = per_layer_head_dim[i];
+      }
+      if (g_log.enabled) {
+        Log("info", "DefaultKeyValueCache: Detected per-layer head_dim variation across " +
+                        std::to_string(layer_count_) + " KV cache layers");
+      }
+
+      // Create per-layer empty past tensors since head_dim varies across layers
+      empty_pasts_.resize(layer_count_);
+      for (int i = 0; i < layer_count_; ++i) {
+        std::array<int64_t, 4> empty_shape = layer_shapes_[i];
+        empty_shape[2] = 0;  // sequence length = 0 for empty past
+        empty_pasts_[i] = OrtValue::CreateTensor(Allocator(), empty_shape, type_);
+      }
+    }
+  }
+
   if (state_.params_->use_graph_capture && !past_present_share_buffer_) {
     // share buffer is a precondition for graph capture
     throw std::runtime_error("Graph capture is not supported with past_present_share_buffer set to false.");
@@ -251,6 +294,13 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state)
     }
   } else if (past_present_share_buffer_) {
     shape_[2] = state_.params_->search.max_length;
+
+    // If per-layer shapes exist (from head_dim auto-detection), update their sequence dim too
+    if (!layer_shapes_.empty()) {
+      for (int i = 0; i < layer_count_; ++i) {
+        layer_shapes_[i][2] = state_.params_->search.max_length;
+      }
+    }
   }
 
   try {
@@ -286,7 +336,12 @@ void DefaultKeyValueCache::Add() {
   output_index_ = state_.outputs_.size();
 
   for (int i = 0; i < layer_count_ * 2; ++i) {
-    state_.inputs_.push_back(empty_past_.get());  // Set empty past here, Update() takes care of the rest
+    // Use per-layer empty past when head_dim varies across layers
+    if (!empty_pasts_.empty()) {
+      state_.inputs_.push_back(empty_pasts_[i / 2].get());
+    } else {
+      state_.inputs_.push_back(empty_past_.get());
+    }
     state_.input_names_.push_back(input_name_strings_[i].c_str());
     state_.outputs_.push_back(presents_[i].get());
     state_.output_names_.push_back(output_name_strings_[i].c_str());
@@ -321,7 +376,8 @@ void DefaultKeyValueCache::Update(DeviceSpan<int32_t> beam_indices, int total_le
     for (int layer_idx = 0; layer_idx < layer_count_; ++layer_idx) {
       std::array<int64_t, 4> current_shape = layer_shapes_[layer_idx];
       const int max_cache_length = static_cast<int>(layer_shapes_[layer_idx][2]);
-      current_shape[2] = std::min(total_length, max_cache_length);
+      // If max_cache_length is 0 (unconstrained), use total_length directly
+      current_shape[2] = (max_cache_length > 0) ? std::min(total_length, max_cache_length) : total_length;
 
       // Key tensor
       presents_[layer_idx * 2] = OrtValue::CreateTensor(Allocator(), current_shape, type_);
@@ -354,7 +410,11 @@ void DefaultKeyValueCache::RewindTo(size_t index) {
   if (index == 0) {
     for (int i = 0; i < layer_count_ * 2; i++) {
       pasts_[i] = nullptr;
-      state_.inputs_[input_index_ + i] = empty_past_.get();
+      if (!empty_pasts_.empty()) {
+        state_.inputs_[input_index_ + i] = empty_pasts_[i / 2].get();
+      } else {
+        state_.inputs_[input_index_ + i] = empty_past_.get();
+      }
     }
   } else if (type_ == Ort::TypeToTensorType<float>) {
     RewindPastTensorsTo<float>(index);

src/models/kv_cache.h

@@ -104,6 +104,7 @@ struct DefaultKeyValueCache : KeyValueCache {
   std::vector<std::array<int64_t, 4>> layer_shapes_;
 
   std::unique_ptr<OrtValue> empty_past_;
+  std::vector<std::unique_ptr<OrtValue>> empty_pasts_;  // Per-layer empty past tensors (for varying head_dim)
   std::vector<std::unique_ptr<OrtValue>> pasts_, presents_;
   std::vector<std::string> input_name_strings_, output_name_strings_;
 };

src/models/model.cpp

@@ -833,8 +833,11 @@ std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, std::unique_ptr<Config> conf
     return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, false);
   if (ModelType::IsPipe(config->model.type))
     return std::make_shared<DecoderOnlyPipelineModel>(std::move(config), ort_env);
-  if (ModelType::IsMMM(config->model.type))
-    return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, true);
+  if (ModelType::IsMMM(config->model.type)) {
+    // Auto-detect speech support: if the config has a speech model filename, enable it
+    bool has_speech = !config->model.speech.filename.empty();
+    return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, has_speech);
+  }
   if (config->model.type == "marian-ssru")
     return std::make_shared<MarianModel>(std::move(config), ort_env);
 
@@ -919,6 +922,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess
           {"whisper", Processor::Create<WhisperProcessor>},
           {"phi4mm", Processor::Create<PhiMultiModalProcessor>},
           {"gemma3", Processor::Create<GemmaImageProcessor>},
+          {"gemma4", Processor::Create<Gemma4MultiModalProcessor>},
           {"mistral3", Processor::Create<Mistral3ImageProcessor>},
           {"fara", Processor::Create<QwenImageProcessor>},
           {"qwen2_5_vl", Processor::Create<QwenImageProcessor>},

src/models/model.h

@@ -11,6 +11,7 @@
 #include "whisper_processor.h"
 #include "phi_multimodal_processor.h"
 #include "gemma_image_processor.h"
+#include "gemma4_multimodal_processor.h"
 #include "adapters.h"
 #include "extra_outputs.h"
 

src/models/model_type.h

@@ -15,7 +15,7 @@ namespace Generators {
 struct ModelType {
   inline static bool IsLLM(const std::string& model_type) {
     // Large-language model (LLM)
-    static constexpr std::array<std::string_view, 21> LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gpt2", "gptoss", "granite", "internlm2", "llama", "mistral", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen3", "smollm3"};
+    static constexpr std::array<std::string_view, 22> LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma4_text", "gpt2", "gptoss", "granite", "internlm2", "llama", "mistral", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen3", "smollm3"};
     return std::find(LLM.begin(), LLM.end(), model_type) != LLM.end();
   }
 
@@ -49,7 +49,7 @@ struct ModelType {
 
   inline static bool IsMMM(const std::string& model_type) {
     // Multi-modal model (MMM)
-    static constexpr std::array<std::string_view, 1> MMM = {"phi4mm"};
+    static constexpr std::array<std::string_view, 2> MMM = {"gemma4", "phi4mm"};
     return std::find(MMM.begin(), MMM.end(), model_type) != MMM.end();
   }
 

src/models/multi_modal.cpp

@@ -579,9 +579,11 @@ SpeechState::SpeechState(const MultiModalLanguageModel& model, const GeneratorPa
 void SpeechState::SetExtraInputs(const std::vector<ExtraInput>& extra_inputs, const int64_t num_audio_tokens) {
   num_audio_tokens_ = num_audio_tokens;
 
-  audio_features_ = std::make_unique<MultiModalFeatures>(*this, MultiModalFeatures::Mode::Output,  // Model output
+  // Allocate 3D [batch, num_audio_tokens, hidden_size] matching the speech ONNX model's
+  // output rank. Will be reshaped to 2D before passing to the embedding model.
+  audio_features_ = std::make_unique<MultiModalFeatures>(*this, MultiModalFeatures::Mode::Output,
                                                          model_.config_->model.speech.outputs.audio_features,
-                                                         -1, num_audio_tokens_);
+                                                         params_->BatchBeamSize(), num_audio_tokens_);
   audio_features_->Add();
   extra_inputs_.Add(extra_inputs, model_.speech_session_->GetInputNames());
 }
@@ -616,13 +618,21 @@ void EmbeddingState::SetExtraInputs(const int64_t num_images, const int64_t num_
                                                            model_.config_->model.embedding.inputs.audio_features,
                                                            -1, num_audio_tokens_);
     audio_features_->Add();
+  } else if (model_.session_info_.HasInput(model_.config_->model.embedding.inputs.audio_features)) {
+    // No speech session, but embedding model requires audio_features — provide empty tensor with shape (0, hidden_size)
+    audio_features_ = std::make_unique<MultiModalFeatures>(*this, MultiModalFeatures::Mode::Input,
+                                                           model_.config_->model.embedding.inputs.audio_features,
+                                                           -1, 0);
+    audio_features_->Add();
+    // Pre-allocate an empty tensor since there's no speech session to provide one via ReuseFeaturesBuffer
+    audio_features_->AllocateEmptyFeatures();
   }
 }
 
 void EmbeddingState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, bool is_prompt) {
   input_ids_.Update(next_tokens);
   if (model_.vision_session_) image_features_->Update(is_prompt);
-  if (model_.speech_session_) audio_features_->Update(is_prompt);
+  if (audio_features_) audio_features_->Update(is_prompt);
 }
 
 DeviceSpan<float> EmbeddingState::Run(int current_length, DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices) {
@@ -639,6 +649,13 @@ DecoderState::DecoderState(const MultiModalLanguageModel& model, DeviceSpan<int3
       position_inputs_{CreatePositionInputs(*this, sequence_lengths, model_.config_->model.decoder.inputs.attention_mask)},
       recurrent_state_{CreateRecurrentState(*this)} {
   inputs_embeds_.Add();
+
+  // Some multimodal decoders (e.g., Gemma4) require input_ids alongside inputs_embeds
+  if (model_.session_info_.HasInput(model_.config_->model.decoder.inputs.input_ids)) {
+    decoder_input_ids_ = std::make_unique<DefaultInputIDs>(*this);
+    decoder_input_ids_->Add();
+  }
+
   position_inputs_->Add();
   logits_.Add();
   kv_cache_.Add();
@@ -659,6 +676,7 @@ DeviceSpan<float> DecoderState::Run(int current_length, DeviceSpan<int32_t>& nex
 void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int total_length, DeviceSpan<int32_t> beam_indices) {
   int batch_size = static_cast<int>(inputs_embeds_.GetShape()[0]);
   size_t new_length = next_tokens.size() / batch_size;
+  if (decoder_input_ids_) decoder_input_ids_->Update(next_tokens);
   position_inputs_->Update(next_tokens, total_length, static_cast<int>(new_length));
   kv_cache_.Update(beam_indices, total_length);
   if (recurrent_state_)
@@ -669,6 +687,7 @@ void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int tot
 
 // Overload for pipeline to call
 void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int total_length, DeviceSpan<int32_t> beam_indices, size_t new_length) {
+  if (decoder_input_ids_) decoder_input_ids_->Update(next_tokens);
   kv_cache_.Update(beam_indices, total_length);
   if (recurrent_state_)
     recurrent_state_->Update();
@@ -756,7 +775,19 @@ DeviceSpan<float> MultiModalPipelineState::Run(int current_length, DeviceSpan<in
     if (vision_state_) {
       embedding_state_->image_features_->ReuseFeaturesBuffer(*vision_state_->image_features_);
     }
-    if (speech_state_) embedding_state_->audio_features_->ReuseFeaturesBuffer(*speech_state_->audio_features_);
+    if (speech_state_ && num_audio_tokens_ > 0) {
+      // Reshape speech output from 3D [B, T, hidden] to 2D [B*T, hidden]
+      // to match embedding model's expected 2D audio_features input rank.
+      auto& speech_shape = speech_state_->audio_features_->GetShape();
+      if (speech_shape.size() == 3) {
+        speech_state_->audio_features_->ReshapeFeatures(
+            {speech_shape[0] * speech_shape[1], speech_shape[2]});
+      }
+      embedding_state_->audio_features_->ReuseFeaturesBuffer(*speech_state_->audio_features_);
+    } else if (embedding_state_->audio_features_) {
+      // No audio: provide empty 2D tensor [0, hidden_size] for the embedding model
+      embedding_state_->audio_features_->AllocateEmptyFeatures();
+    }
     embedding_state_->inputs_embeds_.ReuseEmbeddingsBuffer(decoder_state_->inputs_embeds_);
     embedding_state_->Run(current_length, next_tokens, next_indices);
 

src/models/multi_modal.h

@@ -9,6 +9,7 @@
 #include "extra_inputs.h"
 #include "logits.h"
 #include "kv_cache.h"
+#include "input_ids.h"
 #include "position_inputs.h"
 #include "model_type.h"
 #include "recurrent_state.h"
@@ -145,6 +146,7 @@ struct DecoderState : State {
   const MultiModalLanguageModel& model_;
   Embeddings inputs_embeds_{*this, Embeddings::Mode::Input,  // Model input
                             model_.config_->model.decoder.inputs.embeddings};
+  std::unique_ptr<DefaultInputIDs> decoder_input_ids_;       // Optional model input (e.g., Gemma4 decoder needs input_ids)
   std::unique_ptr<PositionInputs> position_inputs_;  // Model input
   DefaultKeyValueCache kv_cache_{*this};             // Model input
   std::unique_ptr<RecurrentState> recurrent_state_;  // Model input (for hybrid models)

src/models/multi_modal_features.cpp

@@ -19,7 +19,8 @@ MultiModalFeatures::MultiModalFeatures(State& state, MultiModalFeatures::Mode mo
                         : model_.session_info_.GetOutputSymbolicShape(name).size();
 
   // If the model expects 3 dimensions, add a batch dimension
-  if (dims == 3) {
+  // batch_size <= 0 signals "skip batch dim even if model has 3D output"
+  if (dims == 3 && batch_size > 0) {
     shape_.push_back(batch_size);
   }
 
@@ -77,4 +78,28 @@ void MultiModalFeatures::ReuseFeaturesBuffer(MultiModalFeatures& other) {
   state_.inputs_[index_] = other.state_.outputs_[other.index_];
 }
 
+void MultiModalFeatures::AllocateEmptyFeatures() {
+  features_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), shape_, type_);
+  state_.inputs_[index_] = features_.get();
+}
+
+void MultiModalFeatures::ReshapeFeatures(std::vector<int64_t> new_shape) {
+  if (!features_) return;
+  auto old_info = features_->GetTensorTypeAndShapeInfo();
+  int64_t old_count = static_cast<int64_t>(old_info->GetElementCount());
+  int64_t new_count = 1;
+  for (auto d : new_shape) new_count *= d;
+  if (old_count != new_count || old_count == 0) return;
+
+  auto old_features = std::move(features_);
+  features_ = OrtValue::CreateTensor(model_.p_device_->GetAllocator(), new_shape, type_);
+  auto src = ByteWrapTensor(*model_.p_device_, *old_features);
+  auto dst = ByteWrapTensor(*model_.p_device_, *features_);
+  dst.CopyFrom(src);
+  shape_ = std::move(new_shape);
+  if (mode_ == Mode::Output && index_ != ~0U) {
+    state_.outputs_[index_] = features_.get();
+  }
+}
+
 }  // namespace Generators