diff --git a/cmake/deps.txt b/cmake/deps.txt
index a01baa39d0..6920cabc7c 100644
--- a/cmake/deps.txt
+++ b/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
diff --git a/examples/python/common.py b/examples/python/common.py
index d77b83f2bc..9502fb16c6 100644
--- a/examples/python/common.py
+++ b/examples/python/common.py
@@ -4,10 +4,11 @@
 import argparse
 import json
 import os
+from dataclasses import asdict, dataclass
+from typing import Any
+
 import onnxruntime_genai as og
 
-from dataclasses import dataclass, asdict
-from typing import Any
 
 def set_logger(inputs: bool = True, outputs: bool = True) -> None:
     """
@@ -21,6 +22,7 @@ def set_logger(inputs: bool = True, outputs: bool = True) -> None:
     """
     og.set_log_options(enabled=True, model_input_values=inputs, model_output_values=outputs)
 
+
 def register_ep(ep: str, ep_path: str, use_winml: bool) -> None:
     """
     Register execution provider if path is provided or via Windows ML
@@ -42,6 +44,7 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None:
         # Modified from here: https://learn.microsoft.com/en-us/windows/ai/new-windows-ml/tutorial?tabs=python#acquiring-the-model-and-preprocessing
         try:
             import winml
+
             print(winml.register_execution_providers(ort=False, ort_genai=True))
         except ImportError:
             print("WinML not available, using default execution providers")
@@ -53,11 +56,14 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None:
         og.register_execution_provider_library("NvTensorRTRTXExecutionProvider", ep_path)
     else:
         print(f"Warning: EP registration not supported for {ep}")
-        print("Only 'cuda' and 'NvTensorRtRtx' support plug-in libraries. Use Windows ML via '--use_winml' to register EPs.")
+        print(
+            "Only 'cuda' and 'NvTensorRtRtx' support plug-in libraries. Use Windows ML via '--use_winml' to register EPs."
+        )
         return
 
     print(f"Registered {ep} successfully!")
 
+
 def get_config(path: str, ep: str, ep_options: dict[str, str] = {}, search_options: dict[str, int] = {}) -> og.Config:
     """
     Get og.Config object and set EP-specific and search-specific options inside it
@@ -98,6 +104,7 @@ def get_config(path: str, ep: str, ep_options: dict[str, str] = {}, search_optio
     config.overlay(json.dumps({"search": search_options}))
     return config
 
+
 def get_search_options(args: argparse.Namespace):
     """
     Get search options for a generator's params during decoding
@@ -128,7 +135,10 @@ def get_search_options(args: argparse.Namespace):
     search_options["batch_size"] = search_options.get("batch_size", 1)
     return search_options
 
-def apply_chat_template(model_path: str, tokenizer: og.Tokenizer, messages: str, add_generation_prompt: bool, tools: str = "") -> str:
+
+def apply_chat_template(
+    model_path: str, tokenizer: og.Tokenizer, messages: str, add_generation_prompt: bool, tools: str = ""
+) -> str:
     """
     Apply the chat template with various fallback options
 
@@ -151,6 +161,7 @@ def apply_chat_template(model_path: str, tokenizer: og.Tokenizer, messages: str,
     )
     return prompt
 
+
 def get_user_prompt(prompt: str, non_interactive: bool) -> str:
     """
     Get prompt for 'user' role in chat template
@@ -179,6 +190,7 @@ def get_user_prompt(prompt: str, non_interactive: bool) -> str:
 
     return text
 
+
 def get_user_media_paths(media_paths: list[str], non_interactive: bool, media_type: str) -> list[str]:
     """
     Get paths to media for user
@@ -202,7 +214,9 @@ def get_user_media_paths(media_paths: list[str], non_interactive: bool, media_ty
         # If interactive mode is on
         paths = [
             path.strip()
-            for path in input(f"{media_type.capitalize()} Path (comma separated; leave empty if no {media_type}): ").split(",")
+            for path in input(
+                f"{media_type.capitalize()} Path (comma separated; leave empty if no {media_type}): "
+            ).split(",")
         ]
 
     paths = [path for path in paths if path]
@@ -213,6 +227,7 @@ def get_user_media_paths(media_paths: list[str], non_interactive: bool, media_ty
 
     return paths
 
+
 def get_user_images(image_paths: list[str], non_interactive: bool) -> tuple[og.Images, int]:
     """
     Get images for user
@@ -232,6 +247,7 @@ def get_user_images(image_paths: list[str], non_interactive: bool) -> tuple[og.I
     images = og.Images.open(*paths)
     return images, len(paths)
 
+
 def get_user_audios(audio_paths: list[str], non_interactive: bool) -> tuple[og.Audios, int]:
     """
     Get audios for user
@@ -251,6 +267,7 @@ def get_user_audios(audio_paths: list[str], non_interactive: bool) -> tuple[og.A
     audios = og.Audios.open(*paths)
     return audios, len(paths)
 
+
 def get_user_content(model_type: str, num_images: int, num_audios: int, prompt: str) -> str | list[dict[str, str]]:
     """
     Get content for 'user' role in chat template
@@ -284,49 +301,59 @@ 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
 class ToolSchema:
     """
     A class for defining a tool in a JSON schema compatible way
     """
+
     description: str
     type: str
     properties: dict[str, Any]
     required: list[str]
     additionalProperties: bool
 
+
 @dataclass
 class JsonSchema:
     """
     A class for defining a JSON schema for guidance
     """
+
     x_guidance: dict[str, Any]
     type: str
     items: dict[str, list[ToolSchema]]
     minItems: int
 
+
 @dataclass
 class FunctionDefinition:
     """
     A class for defining a function in an OpenAI-compatible way
     """
+
     name: str
     description: str
     parameters: dict[str, Any]
 
+
 @dataclass
 class Tool:
     """
     A class for defining a tool in an OpenAI-compatible way
     """
+
     type: str
     function: FunctionDefinition
 
+
 def tools_to_schemas(tools: list[Tool]) -> list[ToolSchema]:
     """
     Convert a list of tools to a list of tool schemas
@@ -360,6 +387,7 @@ def tools_to_schemas(tools: list[Tool]) -> list[ToolSchema]:
         tool_schemas.append(tool_schema)
     return tool_schemas
 
+
 def get_json_schema(tools: list[Tool], tool_output: bool) -> str:
     """
     Create a JSON schema from a list of tools
@@ -376,6 +404,7 @@ def get_json_schema(tools: list[Tool], tool_output: bool) -> str:
     d = {k.replace("x_guidance", "x-guidance"): v for k, v in asdict(json_schema).items()}
     return json.dumps(d)
 
+
 def get_lark_grammar(
     tools: list[Tool],
     text_output: bool,
@@ -423,6 +452,7 @@ def get_lark_grammar(
 
     return "\n".join(rows)
 
+
 def to_tool(tool_defs: list[dict[str, Any]]) -> list[Tool]:
     """
     Convert a JSON-deserialized object of tools to a list of Tool objects
@@ -443,6 +473,7 @@ def to_tool(tool_defs: list[dict[str, Any]]) -> list[Tool]:
         tools.append(tool)
     return tools
 
+
 def get_guidance(
     response_format: str = "",
     filepath: str = "",
@@ -474,7 +505,7 @@ def get_guidance(
     if tool_output:
         if os.path.exists(filepath):
             # If tools are provided as a file
-            with open(filepath, 'r') as f:
+            with open(filepath) as f:
                 tool_defs = json.load(f)
                 tools = to_tool(tool_defs)
         elif tools_str != "":
@@ -488,14 +519,18 @@ def get_guidance(
             if type(tools[0]) != Tool:
                 tools = to_tool(tools)
         else:
-            raise ValueError("Please provide the list of tools through a file, JSON-serialized string, or a list of tools")
+            raise ValueError(
+                "Please provide the list of tools through a file, JSON-serialized string, or a list of tools"
+            )
 
         assert len(tools) > 0, "Could not obtain a list of tools in memory"
 
     # Create guidance based on user-provided response format
     if response_format in {"text", "lark_grammar"}:
         if response_format == "text":
-            assert text_output and not tool_output, "A response format of 'text' requires text_output = True and tool_output = False"
+            assert text_output and not tool_output, (
+                "A response format of 'text' requires text_output = True and tool_output = False"
+            )
 
         guidance_type = "lark_grammar"
         guidance_data = get_lark_grammar(
@@ -506,7 +541,9 @@ def get_guidance(
             tool_call_end=tool_call_end,
         )
     elif response_format in {"json_schema", "json_object"}:
-        assert tool_output and not text_output, "A response format of 'json_schema' or 'json_object' requires text_output = False and tool_output = True"
+        assert tool_output and not text_output, (
+            "A response format of 'json_schema' or 'json_object' requires text_output = False and tool_output = True"
+        )
 
         guidance_type = "json_schema"
         guidance_data = get_json_schema(tools=tools, tool_output=tool_output)
@@ -515,6 +552,7 @@ def get_guidance(
 
     return guidance_type, guidance_data, json.dumps([asdict(tool) for tool in tools])
 
+
 def get_generator_params_args(parser: argparse.ArgumentParser) -> None:
     """
     Add an argument group for the generator params
@@ -525,16 +563,34 @@ def get_generator_params_args(parser: argparse.ArgumentParser) -> None:
         None
     """
     generator_params = parser.add_argument_group("Generator Params")
-    generator_params.add_argument('-c', '--chunk_size', type=int, default=0, help="Chunk size for prefill chunking during context processing (default: 0 = disabled, >0 = enabled)")
-    generator_params.add_argument('-s', '--do_sample', action='store_true', help='Do random sampling. When false, greedy or beam search are used to generate the output. Defaults to false')
-    generator_params.add_argument('-i', '--min_length', type=int, help='Min number of tokens to generate including the prompt')
-    generator_params.add_argument('-l', '--max_length', type=int, help='Max number of tokens to generate including the prompt')
-    generator_params.add_argument('-b', '--num_beams', type=int, default=1, help='Number of beams to create')
-    generator_params.add_argument('-rs', '--num_return_sequences', type=int, default=1, help='Number of return sequences to produce')
-    generator_params.add_argument('-r', '--repetition_penalty', type=float, help='Repetition penalty to sample with')
-    generator_params.add_argument('-t', '--temperature', type=float, help='Temperature to sample with')
-    generator_params.add_argument('-k', '--top_k', type=int, help='Top k tokens to sample from')
-    generator_params.add_argument('-p', '--top_p', type=float, help='Top p probability to sample with')
+    generator_params.add_argument(
+        "-c",
+        "--chunk_size",
+        type=int,
+        default=0,
+        help="Chunk size for prefill chunking during context processing (default: 0 = disabled, >0 = enabled)",
+    )
+    generator_params.add_argument(
+        "-s",
+        "--do_sample",
+        action="store_true",
+        help="Do random sampling. When false, greedy or beam search are used to generate the output. Defaults to false",
+    )
+    generator_params.add_argument(
+        "-i", "--min_length", type=int, help="Min number of tokens to generate including the prompt"
+    )
+    generator_params.add_argument(
+        "-l", "--max_length", type=int, help="Max number of tokens to generate including the prompt"
+    )
+    generator_params.add_argument("-b", "--num_beams", type=int, default=1, help="Number of beams to create")
+    generator_params.add_argument(
+        "-rs", "--num_return_sequences", type=int, default=1, help="Number of return sequences to produce"
+    )
+    generator_params.add_argument("-r", "--repetition_penalty", type=float, help="Repetition penalty to sample with")
+    generator_params.add_argument("-t", "--temperature", type=float, help="Temperature to sample with")
+    generator_params.add_argument("-k", "--top_k", type=int, help="Top k tokens to sample from")
+    generator_params.add_argument("-p", "--top_p", type=float, help="Top p probability to sample with")
+
 
 def get_guidance_args(parser: argparse.ArgumentParser) -> None:
     """
@@ -546,9 +602,38 @@ def get_guidance_args(parser: argparse.ArgumentParser) -> None:
         None
     """
     guidance = parser.add_argument_group("Guidance Arguments")
-    guidance.add_argument('-rf', '--response_format', type=str, default="", choices=["", "text", "json_object", "json_schema", "lark_grammar"], help='Provide response format for the model')
-    guidance.add_argument('-tf', '--tools_file', type=str, default="", help='Path to file containing list of OpenAI-compatible tool definitions. Ex: test/test_models/tool-definitions/weather.json')
-    guidance.add_argument('-text', '--text_output', action='store_true', default=False, help='Produce a text response in the output')
-    guidance.add_argument('-tool', '--tool_output', action='store_true', default=False, help='Produce a tool call in the output')
-    guidance.add_argument('-tcs', '--tool_call_start', type=str, default="", help='String representation of tool call start (ex: <|tool_call|>). Needs to be marked as special in tokenizer.json for guidance to work.')
-    guidance.add_argument('-tce', '--tool_call_end', type=str, default="", help='String representation of tool call end (ex: <|/tool_call|>). Needs to be marked as special in tokenizer.json for guidance to work.')
+    guidance.add_argument(
+        "-rf",
+        "--response_format",
+        type=str,
+        default="",
+        choices=["", "text", "json_object", "json_schema", "lark_grammar"],
+        help="Provide response format for the model",
+    )
+    guidance.add_argument(
+        "-tf",
+        "--tools_file",
+        type=str,
+        default="",
+        help="Path to file containing list of OpenAI-compatible tool definitions. Ex: test/test_models/tool-definitions/weather.json",
+    )
+    guidance.add_argument(
+        "-text", "--text_output", action="store_true", default=False, help="Produce a text response in the output"
+    )
+    guidance.add_argument(
+        "-tool", "--tool_output", action="store_true", default=False, help="Produce a tool call in the output"
+    )
+    guidance.add_argument(
+        "-tcs",
+        "--tool_call_start",
+        type=str,
+        default="",
+        help="String representation of tool call start (ex: <|tool_call|>). Needs to be marked as special in tokenizer.json for guidance to work.",
+    )
+    guidance.add_argument(
+        "-tce",
+        "--tool_call_end",
+        type=str,
+        default="",
+        help="String representation of tool call end (ex: <|/tool_call|>). Needs to be marked as special in tokenizer.json for guidance to work.",
+    )
diff --git a/src/config.cpp b/src/config.cpp
index 0fee4f703c..cab5ff2800 100644
--- a/src/config.cpp
+++ b/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") {
diff --git a/src/config.h b/src/config.h
index 6b98875562..80fd9cba2f 100644
--- a/src/config.h
+++ b/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";
 
@@ -127,8 +128,10 @@ struct Config {
     int sep_token_id{};             // The id of the separation token.
     int decoder_start_token_id{};   // If an encoder-decoder model starts decoding with a different token than bos, the id of that token.
 
-    // Qwen2.5-VL specific token IDs
+    // Multimodal token IDs (used by Qwen-VL, Gemma4, and other VLM/MMM models)
     int image_token_id{};
+    int audio_token_id{};
+    int boa_token_id{};  // Beginning-of-audio 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
diff --git a/src/models/gemma4_multimodal_processor.cpp b/src/models/gemma4_multimodal_processor.cpp
new file mode 100644
index 0000000000..8688a52dd0
--- /dev/null
+++ b/src/models/gemma4_multimodal_processor.cpp
@@ -0,0 +1,387 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "../generators.h"
+#include "model.h"
+
+namespace Generators {
+
+namespace {
+
+// Simple literal string count (no regex overhead for fixed tokens)
+size_t CountOccurrences(const std::string& text, const std::string& token) {
+  size_t count = 0;
+  size_t pos = 0;
+  while ((pos = text.find(token, pos)) != std::string::npos) {
+    ++count;
+    pos += token.size();
+  }
+  return count;
+}
+
+// Replace all occurrences of a literal string (avoids std::regex compilation cost)
+void ReplaceAll(std::string& text, const std::string& from, const std::string& to) {
+  size_t pos = 0;
+  while ((pos = text.find(from, pos)) != std::string::npos) {
+    text.replace(pos, from.size(), to);
+    pos += to.size();
+  }
+}
+
+// Expand image and audio placeholder tokens in the prompt, then encode to input_ids.
+// Returns (input_ids, token_type_ids, num_img_tokens).
+std::tuple<std::unique_ptr<OrtValue>, std::unique_ptr<OrtValue>, std::unique_ptr<OrtValue>>
+ProcessGemma4Prompt(const Generators::Tokenizer& tokenizer, const std::string& prompt,
+                    OrtxTensor* pixel_values, Ort::Allocator& allocator,
+                    size_t vision_soft_tokens_per_image,
+                    int64_t num_audio_tokens = 0) {
+  constexpr char boi_token[] = "<|image>";
+  constexpr char image_token[] = "<|image|>";
+  constexpr char eoi_token[] = "<image|>";
+  constexpr size_t boi_token_len = sizeof(boi_token) - 1;
+  constexpr size_t image_token_len = sizeof(image_token) - 1;
+
+  int64_t num_images{};
+  if (pixel_values) {
+    const float* pixel_values_data{};
+    const int64_t* pixel_values_shape{};
+    size_t pixel_values_num_dims;
+    CheckResult(OrtxGetTensorData(pixel_values, reinterpret_cast<const void**>(&pixel_values_data),
+                                  &pixel_values_shape, &pixel_values_num_dims));
+    // 3D: [batch/num_images, num_patches, patch_dim] → shape[0] is num_images
+    // 2D: [num_patches, patch_dim] → single image (no batch dim)
+    if (pixel_values_num_dims == 3) {
+      num_images = pixel_values_shape[0];
+    } else if (pixel_values_num_dims == 2) {
+      num_images = 1;
+    } else {
+      throw std::runtime_error("pixel_values has unexpected rank " + std::to_string(pixel_values_num_dims) +
+                               ". Expected 2 (num_patches, patch_dim) or 3 (batch, num_patches, patch_dim).");
+    }
+  }
+
+  std::string text = prompt;
+  if (num_images > 0) {
+    // Count existing boi tokens in prompt
+    auto existing_boi_count = static_cast<int64_t>(CountOccurrences(text, boi_token));
+
+    // The chat template may insert <|image|> (image_token) instead of <|image> (boi_token).
+    // If we find standalone <|image|> tokens that aren't part of <|image>, treat each as one image.
+    if (existing_boi_count == 0) {
+      auto image_token_count = static_cast<int64_t>(CountOccurrences(text, image_token));
+      if (image_token_count > 0 && image_token_count == num_images) {
+        // Replace each standalone <|image|> with <|image> so the expansion logic works
+        ReplaceAll(text, image_token, boi_token);
+        existing_boi_count = image_token_count;
+      }
+    }
+
+    if (existing_boi_count == 0) {
+      // No image tokens in prompt — auto-insert them before the text
+      std::string prefix;
+      prefix.reserve(static_cast<size_t>(num_images) * (boi_token_len + 1));
+      for (int64_t i = 0; i < num_images; ++i) {
+        prefix += boi_token;
+        if (i < num_images - 1) prefix += ' ';
+      }
+      text = prefix + (text.empty() ? "" : " ") + text;
+    }
+  }
+
+  // Count and validate boi tokens using simple string search
+  const auto boi_count = static_cast<int64_t>(CountOccurrences(text, boi_token));
+  if (num_images != boi_count) {
+    throw std::runtime_error("Prompt contained " + std::to_string(boi_count) + " image tokens but received " +
+                             std::to_string(num_images) + " images.");
+  }
+
+  // Build the expanded image token sequence with pre-allocated buffer
+  std::string image_tokens_expanded;
+  image_tokens_expanded.reserve(vision_soft_tokens_per_image * image_token_len);
+  for (size_t i = 0; i < vision_soft_tokens_per_image; ++i) {
+    image_tokens_expanded += image_token;
+  }
+  const std::string full_image_sequence = "\n\n" + std::string(boi_token) + image_tokens_expanded + eoi_token + "\n\n";
+  ReplaceAll(text, boi_token, full_image_sequence);
+
+  // Expand audio tokens: replace single <|audio|> from chat template with N audio soft tokens.
+  // Currently only single-clip audio is supported. Multi-audio would require per-clip
+  // token counts from the speech encoder, which batch mel extraction doesn't provide.
+  if (num_audio_tokens > 0) {
+    constexpr char boa_token[] = "<|audio>";
+    constexpr char audio_token[] = "<|audio|>";
+    constexpr char eoa_token[] = "<audio|>";
+
+    // Validate: only single audio clip is supported
+    auto audio_marker_count = CountOccurrences(text, audio_token);
+    if (audio_marker_count > 1) {
+      throw std::runtime_error("Gemma4 audio processing currently supports only 1 audio clip per prompt, but found " +
+                               std::to_string(audio_marker_count) + " audio markers in the prompt.");
+    }
+
+    std::string audio_tokens_expanded;
+    audio_tokens_expanded.reserve(static_cast<size_t>(num_audio_tokens) * (sizeof(audio_token) - 1));
+    for (int64_t i = 0; i < num_audio_tokens; ++i) {
+      audio_tokens_expanded += audio_token;
+    }
+    const std::string full_audio_sequence = "\n\n" + std::string(boa_token) + audio_tokens_expanded + eoa_token + "\n\n";
+
+    // Chat template inserts <|audio|> per audio clip — replace with expanded sequence
+    if (audio_marker_count == 1) {
+      auto pos = text.find(audio_token);
+      if (pos != std::string::npos) {
+        text.replace(pos, sizeof(audio_token) - 1, full_audio_sequence);
+      }
+    } else {
+      // No audio marker in prompt — look for <|audio> (boa_token)
+      auto boa_count = CountOccurrences(text, boa_token);
+      if (boa_count > 0) {
+        ReplaceAll(text, boa_token, full_audio_sequence);
+      } else {
+        // No audio tokens at all — append before the text
+        text = full_audio_sequence + text;
+      }
+    }
+  }
+
+  const std::vector<int32_t> input_ids = tokenizer.Encode(text.c_str());
+  const auto seq_len = static_cast<int64_t>(input_ids.size());
+
+  auto input_ids_value = OrtValue::CreateTensor<int32_t>(allocator, std::vector<int64_t>{1, seq_len});
+  std::copy(input_ids.begin(), input_ids.end(), input_ids_value->GetTensorMutableData<int32_t>());
+
+  auto token_type_ids = OrtValue::CreateTensor<int32_t>(allocator, std::vector<int64_t>{1, seq_len});
+  const auto image_token_id = tokenizer.TokenToTokenId(image_token);
+  auto* token_type_data = token_type_ids->GetTensorMutableData<int32_t>();
+  for (size_t i = 0; i < input_ids.size(); ++i) {
+    token_type_data[i] = (input_ids[i] == image_token_id) ? 1 : 0;
+  }
+
+  auto num_img_tokens = OrtValue::CreateTensor<int64_t>(allocator, std::vector<int64_t>{1});
+  num_img_tokens->GetTensorMutableData<int64_t>()[0] = static_cast<int64_t>(vision_soft_tokens_per_image);
+
+  return {std::move(input_ids_value), std::move(token_type_ids), std::move(num_img_tokens)};
+}
+
+}  // namespace
+
+Gemma4MultiModalProcessor::Gemma4MultiModalProcessor(Config& config, const SessionInfo& session_info)
+    : pixel_values_type_{session_info.GetInputDataType(config.model.vision.inputs.pixel_values)} {
+  // Query pixel_position_ids type (int32 or int64) if the vision model has this input
+  if (session_info.HasInput(config.model.vision.inputs.pixel_position_ids)) {
+    pixel_position_ids_type_ = session_info.GetInputDataType(config.model.vision.inputs.pixel_position_ids);
+  }
+  const auto image_processor_config = (config.config_path / fs::path(config.model.vision.config_filename)).string();
+  CheckResult(OrtxCreateProcessor(image_processor_.ToBeAssigned(), image_processor_config.c_str()));
+
+  config.AddMapping(std::string(Config::Defaults::InputIdsName), config.model.embedding.inputs.input_ids);
+  config.AddMapping(std::string(Config::Defaults::PixelValuesName), config.model.vision.inputs.pixel_values);
+  config.AddMapping(std::string(Config::Defaults::PixelPositionIdsName), config.model.vision.inputs.pixel_position_ids);
+
+  // Initialize speech/audio processor if config is present
+  if (!config.model.speech.config_filename.empty()) {
+    auto speech_config_path = config.config_path / fs::path(config.model.speech.config_filename);
+    if (fs::exists(speech_config_path)) {
+      has_speech_ = true;
+      audio_features_type_ = session_info.GetInputDataType(config.model.speech.inputs.audio_embeds);
+      CheckResult(OrtxCreateSpeechFeatureExtractor(audio_processor_.ToBeAssigned(), speech_config_path.string().c_str()));
+
+      config.AddMapping(std::string(Config::Defaults::AudioEmbedsName), config.model.speech.inputs.audio_embeds);
+      config.AddMapping(std::string(Config::Defaults::AudioAttentionMaskName), config.model.speech.inputs.attention_mask);
+      config.AddMapping(std::string(Config::Defaults::AudioSizesName), config.model.speech.inputs.audio_sizes);
+    } else if (!config.model.speech.filename.empty()) {
+      // Speech model is configured but the preprocessing config file is missing on disk
+      throw std::runtime_error("Speech model is configured (speech.filename=" + config.model.speech.filename +
+                               ") but the audio processor config file was not found at: " +
+                               speech_config_path.string());
+    }
+  }
+}
+
+std::unique_ptr<NamedTensors> Gemma4MultiModalProcessor::Process(const Tokenizer& tokenizer, const Payload& payload) const {
+  Ort::Allocator& allocator{Ort::Allocator::GetWithDefaultOptions()};
+  auto named_tensors = std::make_unique<NamedTensors>();
+
+  // Text-only path: no images and no audio
+  if (!payload.images && !payload.audios) {
+    auto [input_ids, token_type_ids, num_img_tokens] =
+        ProcessGemma4Prompt(tokenizer, std::string(payload.prompt), nullptr, allocator, vision_soft_tokens_per_image_);
+    named_tensors->emplace(Config::Defaults::InputIdsName, std::make_shared<Tensor>(std::move(input_ids)));
+    return named_tensors;
+  }
+
+  // Process images if present
+  ort_extensions::OrtxObjectPtr<OrtxTensorResult> image_result;
+  OrtxTensor* pixel_values = nullptr;
+  OrtxTensor* pixel_position_ids = nullptr;
+  size_t actual_soft_tokens = vision_soft_tokens_per_image_;
+  if (payload.images) {
+    CheckResult(OrtxImagePreProcess(image_processor_.get(), payload.images->images_.get(), image_result.ToBeAssigned()));
+    CheckResult(OrtxTensorResultGetAt(image_result.get(), 0, &pixel_values));
+
+    // pixel_position_ids is the second output from the Gemma4 image preprocessor
+    OrtxTensor* temp_tensor = nullptr;
+    if (OrtxTensorResultGetAt(image_result.get(), 1, &temp_tensor) == kOrtxOK && temp_tensor) {
+      pixel_position_ids = temp_tensor;
+    }
+
+    // num_soft_tokens is the third output — the actual number of vision tokens after pooling
+    OrtxTensor* num_soft_tokens_tensor = nullptr;
+    if (OrtxTensorResultGetAt(image_result.get(), 2, &num_soft_tokens_tensor) == kOrtxOK && num_soft_tokens_tensor) {
+      const int64_t* nst_data{};
+      const int64_t* nst_shape{};
+      size_t nst_dims;
+      CheckResult(OrtxGetTensorData(num_soft_tokens_tensor, reinterpret_cast<const void**>(&nst_data),
+                                    &nst_shape, &nst_dims));
+      if (nst_data && nst_data[0] > 0) {
+        actual_soft_tokens = static_cast<size_t>(nst_data[0]);
+      }
+    }
+  }
+
+  // Process audio FIRST to compute num_audio_tokens (needed for prompt token expansion).
+  // Currently only single-clip audio is supported per prompt.
+  int64_t num_audio_tokens = 0;
+  if (payload.audios && !has_speech_) {
+    throw std::runtime_error(
+        "Audio input was provided but audio/speech support is not configured. "
+        "Ensure the genai_config.json has a 'speech' section with both 'filename' and 'config_filename'.");
+  }
+  if (payload.audios && has_speech_) {
+    ort_extensions::OrtxObjectPtr<OrtxTensorResult> audio_result;
+    CheckResult(OrtxFeatureExtraction(audio_processor_.get(), payload.audios->audios_.get(), audio_result.ToBeAssigned()));
+
+    OrtxTensor* audio_features = nullptr;
+    CheckResult(OrtxTensorResultGetAt(audio_result.get(), 0, &audio_features));
+
+    EmplaceProcessedTensor(*named_tensors, Config::Defaults::AudioEmbedsName, audio_features, audio_features_type_, allocator);
+
+    // Create input_features_mask: all-True for single-clip inference (no padding)
+    // Shape matches audio features: [batch, time] bool
+    const float* audio_data{};
+    const int64_t* audio_shape{};
+    size_t audio_dims;
+    CheckResult(OrtxGetTensorData(audio_features, reinterpret_cast<const void**>(&audio_data),
+                                  &audio_shape, &audio_dims));
+    int64_t time_dim = (audio_dims == 3) ? audio_shape[1] : audio_shape[0];
+    int64_t batch_dim = (audio_dims == 3) ? audio_shape[0] : 1;
+    if (batch_dim > 1) {
+      throw std::runtime_error(
+          "Gemma4 audio processing currently supports only 1 audio clip per prompt, "
+          "but received " +
+          std::to_string(batch_dim) + " clips.");
+    }
+    {
+      auto mask = OrtValue::CreateTensor<bool>(allocator, std::vector<int64_t>{batch_dim, time_dim});
+      std::fill_n(mask->GetTensorMutableData<bool>(), batch_dim * time_dim, true);
+      named_tensors->emplace(std::string(Config::Defaults::AudioAttentionMaskName),
+                             std::make_shared<Tensor>(std::move(mask)));
+    }
+
+    // Compute audio_sizes: the speech encoder uses 2-stage Conv2d with stride=2 each
+    int64_t t_after_1 = (time_dim - 1) / 2 + 1;
+    int64_t t_after_2 = (t_after_1 - 1) / 2 + 1;
+    num_audio_tokens = t_after_2;
+    std::array<int64_t, 1> audio_sizes_shape = {1};
+    auto audio_sizes = OrtValue::CreateTensor<int64_t>(allocator, audio_sizes_shape);
+    audio_sizes->GetTensorMutableData<int64_t>()[0] = num_audio_tokens;
+    named_tensors->emplace(std::string(Config::Defaults::AudioSizesName),
+                           std::make_shared<Tensor>(std::move(audio_sizes)));
+  }
+
+  // Process prompt: expand image and audio tokens, then encode
+  auto [input_ids, token_type_ids, num_img_tokens] =
+      ProcessGemma4Prompt(tokenizer, std::string(payload.prompt), pixel_values, allocator, actual_soft_tokens, num_audio_tokens);
+  named_tensors->emplace(std::string(Config::Defaults::InputIdsName), std::make_shared<Tensor>(std::move(input_ids)));
+  named_tensors->emplace(std::string(Config::Defaults::TokenTypeIdsName), std::make_shared<Tensor>(std::move(token_type_ids)));
+
+  if (payload.images) {
+    // The Gemma4ImageTransform pads pixel_values and position_ids to max_patches.
+    // The vision ONNX model expects the actual (unpadded) number of patches.
+    // Trim the tensors to actual_patches = actual_soft_tokens * pooling_kernel_size².
+    constexpr int64_t kPoolingKernelSize = 3;
+    const int64_t actual_patches = static_cast<int64_t>(actual_soft_tokens) * kPoolingKernelSize * kPoolingKernelSize;
+
+    // Get padded pixel_values shape
+    const float* pv_data{};
+    const int64_t* pv_shape{};
+    size_t pv_dims;
+    CheckResult(OrtxGetTensorData(pixel_values, reinterpret_cast<const void**>(&pv_data), &pv_shape, &pv_dims));
+
+    // Determine the patches dimension and patch_dim based on tensor rank
+    // 2D: (num_patches, patch_dim) or 3D: (batch, num_patches, patch_dim)
+    const int64_t num_padded_patches = (pv_dims == 3) ? pv_shape[1] : pv_shape[0];
+    const int64_t patch_dim = (pv_dims == 3) ? pv_shape[2] : pv_shape[1];
+
+    if (actual_patches < num_padded_patches) {
+      // Trim: copy only the first actual_patches from the padded tensor.
+      // For 3D [batch, patches, dim], copy batch * actual_patches * dim elements.
+      const int64_t batch = (pv_dims == 3) ? pv_shape[0] : 1;
+      auto trimmed_shape = (pv_dims == 3) ? std::vector<int64_t>{batch, actual_patches, patch_dim}
+                                          : std::vector<int64_t>{actual_patches, patch_dim};
+
+      // Respect the model's pixel_values type (float, fp16, or bf16)
+      auto trimmed_pv = OrtValue::CreateTensor(allocator, trimmed_shape, pixel_values_type_);
+      const size_t elem_size = (pixel_values_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) ? 4 : 2;  // float=4, fp16/bf16=2
+      // For 3D with batch > 1, copy each batch slice separately (padded stride differs from trimmed stride).
+      // Use raw byte pointers since the type may be float, fp16, or bf16.
+      auto* dst = static_cast<uint8_t*>(trimmed_pv->GetTensorMutableRawData());
+      const auto* src = reinterpret_cast<const uint8_t*>(pv_data);
+      const size_t src_stride = static_cast<size_t>(num_padded_patches * patch_dim) * elem_size;
+      const size_t dst_stride = static_cast<size_t>(actual_patches * patch_dim) * elem_size;
+      for (int64_t b = 0; b < batch; ++b) {
+        std::memcpy(dst + b * dst_stride, src + b * src_stride, dst_stride);
+      }
+
+      named_tensors->emplace(std::string(Config::Defaults::PixelValuesName),
+                             std::make_shared<Tensor>(std::move(trimmed_pv)));
+    } else {
+      EmplaceProcessedTensor(*named_tensors, Config::Defaults::PixelValuesName, pixel_values, pixel_values_type_, allocator);
+    }
+
+    named_tensors->emplace(std::string(Config::Defaults::NumImageTokens), std::make_shared<Tensor>(std::move(num_img_tokens)));
+
+    // Trim pixel_position_ids similarly
+    if (pixel_position_ids) {
+      const void* pos_data_raw{};
+      const int64_t* pos_shape{};
+      size_t pos_dims;
+      CheckResult(OrtxGetTensorData(pixel_position_ids, &pos_data_raw, &pos_shape, &pos_dims));
+
+      const int64_t num_padded_pos = (pos_dims == 3) ? pos_shape[1] : pos_shape[0];
+      const int64_t pos_last_dim = (pos_dims == 3) ? pos_shape[2] : pos_shape[1];
+
+      if (actual_patches < num_padded_pos) {
+        // Trim position_ids: for 3D, copy per-batch with correct stride.
+        // Detect the element type from the vision model's input to handle both int32 and int64.
+        const int64_t pos_batch = (pos_dims == 3) ? pos_shape[0] : 1;
+        auto trimmed_pos_shape = (pos_dims == 3) ? std::vector<int64_t>{pos_batch, actual_patches, pos_last_dim}
+                                                 : std::vector<int64_t>{actual_patches, pos_last_dim};
+
+        auto trimmed_pos = OrtValue::CreateTensor(allocator, trimmed_pos_shape, pixel_position_ids_type_);
+        const size_t pos_elem_size = (pixel_position_ids_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) ? 4 : 8;
+        auto* dst = static_cast<uint8_t*>(trimmed_pos->GetTensorMutableRawData());
+        const auto* src = static_cast<const uint8_t*>(pos_data_raw);
+        const size_t src_stride = static_cast<size_t>(num_padded_pos * pos_last_dim) * pos_elem_size;
+        const size_t dst_stride = static_cast<size_t>(actual_patches * pos_last_dim) * pos_elem_size;
+        for (int64_t b = 0; b < pos_batch; ++b) {
+          std::memcpy(dst + b * dst_stride, src + b * src_stride, dst_stride);
+        }
+        named_tensors->emplace(std::string(Config::Defaults::PixelPositionIdsName),
+                               std::make_shared<Tensor>(std::move(trimmed_pos)));
+      } else {
+        if (pixel_position_ids_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) {
+          named_tensors->emplace(std::string(Config::Defaults::PixelPositionIdsName),
+                                 std::make_shared<Tensor>(ProcessTensor<int32_t>(pixel_position_ids, allocator)));
+        } else {
+          named_tensors->emplace(std::string(Config::Defaults::PixelPositionIdsName),
+                                 std::make_shared<Tensor>(ProcessTensor<int64_t>(pixel_position_ids, allocator)));
+        }
+      }
+    }
+  }
+
+  return named_tensors;
+}
+
+}  // namespace Generators
diff --git a/src/models/gemma4_multimodal_processor.h b/src/models/gemma4_multimodal_processor.h
new file mode 100644
index 0000000000..f53e4ec061
--- /dev/null
+++ b/src/models/gemma4_multimodal_processor.h
@@ -0,0 +1,26 @@
+// 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 pixel_position_ids_type_{ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64};
+  ONNXTensorElementDataType audio_features_type_;
+
+  bool has_speech_{false};
+  size_t vision_soft_tokens_per_image_{260};
+};
+
+}  // namespace Generators
diff --git a/src/models/gemma_image_processor.cpp b/src/models/gemma_image_processor.cpp
index 9bd3d711c5..9dbab3feb8 100644
--- a/src/models/gemma_image_processor.cpp
+++ b/src/models/gemma_image_processor.cpp
@@ -109,16 +109,7 @@ std::unique_ptr<NamedTensors> GemmaImageProcessor::Process(const Tokenizer& toke
   named_tensors->emplace(std::string(Config::Defaults::InputIdsName), std::make_shared<Tensor>(std::move(input_ids)));
   named_tensors->emplace(std::string(Config::Defaults::TokenTypeIdsName), std::make_shared<Tensor>(std::move(token_type_ids)));
 
-  if (pixel_values_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
-    named_tensors->emplace(std::string(Config::Defaults::PixelValuesName),
-                           std::make_shared<Tensor>(ProcessTensor<float>(pixel_values, allocator)));
-  } else if (pixel_values_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16) {
-    named_tensors->emplace(std::string(Config::Defaults::PixelValuesName),
-                           std::make_shared<Tensor>(ProcessTensor<Ort::BFloat16_t>(pixel_values, allocator)));
-  } else {
-    named_tensors->emplace(std::string(Config::Defaults::PixelValuesName),
-                           std::make_shared<Tensor>(ProcessTensor<Ort::Float16_t>(pixel_values, allocator)));
-  }
+  EmplaceProcessedTensor(*named_tensors, Config::Defaults::PixelValuesName, pixel_values, pixel_values_type_, allocator);
 
   named_tensors->emplace(std::string(Config::Defaults::NumImageTokens), std::make_shared<Tensor>(std::move(num_img_tokens)));
 
diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp
index 5e0a1be7ba..44de037ef0 100644
--- a/src/models/kv_cache.cpp
+++ b/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.");
@@ -220,12 +263,18 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state)
 
     // Check if we need per-layer allocation for models with alternating attention patterns
     if (!model_.config_->model.decoder.sliding_window->layers.empty()) {
-      // Use per-layer allocation based on sliding window layer indices
-      layer_shapes_.resize(layer_count_);
+      // Use per-layer allocation based on sliding window layer indices.
+      // If layer_shapes_ already exists (from head_dim auto-detection), preserve
+      // the per-layer head_dim values — only update the sequence length dimension.
+      if (layer_shapes_.empty()) {
+        layer_shapes_.resize(layer_count_);
+        for (int layer_idx = 0; layer_idx < layer_count_; ++layer_idx) {
+          layer_shapes_[layer_idx] = shape_;
+        }
+      }
 
-      // Initialize all layers with base shape and max_length
+      // Set all layers to max_length (sequence dim only)
       for (int layer_idx = 0; layer_idx < layer_count_; ++layer_idx) {
-        layer_shapes_[layer_idx] = shape_;
         layer_shapes_[layer_idx][2] = max_length;
       }
 
@@ -251,6 +300,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 +342,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 +382,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 +416,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);
diff --git a/src/models/kv_cache.h b/src/models/kv_cache.h
index 35167b354a..b7cab09f1f 100644
--- a/src/models/kv_cache.h
+++ b/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_;
 };
diff --git a/src/models/model.cpp b/src/models/model.cpp
index 2313012d59..1a7ffcd7f2 100644
--- a/src/models/model.cpp
+++ b/src/models/model.cpp
@@ -833,8 +833,24 @@ 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: require both the speech ONNX model filename
+    // and the preprocessing config to be present. If only one is set, throw
+    // a clear error so misconfigurations don't silently disable audio.
+    bool has_speech_model = !config->model.speech.filename.empty();
+    bool has_speech_config = !config->model.speech.config_filename.empty();
+    if (has_speech_model && !has_speech_config) {
+      throw std::runtime_error(
+          "speech.filename is set but speech.config_filename is missing. "
+          "Both are required for audio support.");
+    }
+    if (!has_speech_model && has_speech_config) {
+      throw std::runtime_error(
+          "speech.config_filename is set but speech.filename is missing. "
+          "Both are required for audio support.");
+    }
+    return std::make_shared<MultiModalLanguageModel>(std::move(config), ort_env, true, has_speech_model);
+  }
   if (config->model.type == "marian-ssru")
     return std::make_shared<MarianModel>(std::move(config), ort_env);
 
@@ -919,6 +935,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>},
diff --git a/src/models/model.h b/src/models/model.h
index 774aa1afbd..cd47f37112 100644
--- a/src/models/model.h
+++ b/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"
 
diff --git a/src/models/model_type.h b/src/models/model_type.h
index 41e1cfcdf1..e1c2e34d15 100644
--- a/src/models/model_type.h
+++ b/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();
   }
 
diff --git a/src/models/multi_modal.cpp b/src/models/multi_modal.cpp
index d30a8a8729..6a628168bd 100644
--- a/src/models/multi_modal.cpp
+++ b/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);
 
diff --git a/src/models/multi_modal.h b/src/models/multi_modal.h
index bd518ceaa5..da60949edb 100644
--- a/src/models/multi_modal.h
+++ b/src/models/multi_modal.h
@@ -145,10 +145,11 @@ struct DecoderState : State {
   const MultiModalLanguageModel& model_;
   Embeddings inputs_embeds_{*this, Embeddings::Mode::Input,  // Model input
                             model_.config_->model.decoder.inputs.embeddings};
-  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)
-  Logits logits_{*this};                             // Model output
+  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)
+  Logits logits_{*this};                                // Model output
 };
 
 struct MultiModalPipelineState : State {
diff --git a/src/models/multi_modal_features.cpp b/src/models/multi_modal_features.cpp
index 2dccd6000f..9577839655 100644
--- a/src/models/multi_modal_features.cpp
+++ b/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,31 @@ void MultiModalFeatures::ReuseFeaturesBuffer(MultiModalFeatures& other) {
   state_.inputs_[index_] = other.state_.outputs_[other.index_];
 }
 
+void MultiModalFeatures::AllocateEmptyFeatures() {
+  // Skip if already allocated (avoids redundant allocation when called from
+  // both EmbeddingState::SetExtraInputs and the pipeline prompt path)
+  if (features_ && state_.inputs_[index_] == features_.get()) return;
+  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
diff --git a/src/models/multi_modal_features.h b/src/models/multi_modal_features.h
index f61c6580f8..bffea8f410 100644
--- a/src/models/multi_modal_features.h
+++ b/src/models/multi_modal_features.h
@@ -19,6 +19,14 @@ struct MultiModalFeatures {
   void Update(bool is_prompt);
   void ReuseFeaturesBuffer(MultiModalFeatures& other);
 
+  // Pre-allocate an empty features tensor for Input mode when no source session provides one.
+  // Used when the embedding model requires an input (e.g., audio_features) but no corresponding
+  // encoder session exists.
+  void AllocateEmptyFeatures();
+
+  // Reshape features tensor in-place (e.g., flatten 3D [B, T, H] to 2D [B*T, H])
+  void ReshapeFeatures(std::vector<int64_t> new_shape);
+
   auto& GetShape() const { return shape_; }
   OrtValue* Get() { return features_.get(); }
 
diff --git a/src/models/position_inputs.cpp b/src/models/position_inputs.cpp
index 78ed9f63fc..7cf3f512c2 100644
--- a/src/models/position_inputs.cpp
+++ b/src/models/position_inputs.cpp
@@ -449,16 +449,16 @@ WindowedPositionInputs::WindowedPositionInputs(State& state)
 
   if (has_posid_input_) {
     position_ids_type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.position_ids);
-    if (position_ids_type_ != Ort::TypeToTensorType<int32_t>)
-      throw std::runtime_error("WindowedPositionInputs only supports int32_t position_ids");
+    if (position_ids_type_ != Ort::TypeToTensorType<int32_t> && position_ids_type_ != Ort::TypeToTensorType<int64_t>)
+      throw std::runtime_error("WindowedPositionInputs only supports int32_t or int64_t position_ids");
 
     position_ids_shape_ = {1, model_.config_->model.decoder.sliding_window->window_size};
   }
 
   if (has_mask_input_) {
     attention_mask_type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
-    if (attention_mask_type_ != Ort::TypeToTensorType<int32_t>)
-      throw std::runtime_error("WindowedPositionInputs only supports int32_t attention_mask");
+    if (attention_mask_type_ != Ort::TypeToTensorType<int32_t> && attention_mask_type_ != Ort::TypeToTensorType<int64_t>)
+      throw std::runtime_error("WindowedPositionInputs only supports int32_t or int64_t attention_mask");
 
     attention_mask_shape_ = {1, model_.config_->model.context_length};
   }
@@ -492,14 +492,20 @@ void WindowedPositionInputs::Update(DeviceSpan<int32_t> next_tokens, int total_l
       // next_tokens -> [0, a, b, c, d, e]
       // window_size = 3, num_windows = 2, pad_token = 0
       // window_index = 0, position_ids_ -> [0, 0, 1]
-      auto* position_ids_data = position_ids_->GetTensorMutableData<int32_t>();
-      for (int i = 0, j = 0; i < position_ids_shape_[1]; i++) {
-        if (next_tokens.Span()[i] == model_.config_->model.pad_token_id) {
-          position_ids_data[i] = 0;
-        } else {
-          position_ids_data[i] = j++;
+      auto fill_first_window = [&](auto* position_ids_data) {
+        using T = std::remove_pointer_t<decltype(position_ids_data)>;
+        for (int i = 0, j = 0; i < position_ids_shape_[1]; i++) {
+          if (next_tokens.Span()[i] == model_.config_->model.pad_token_id) {
+            position_ids_data[i] = T{0};
+          } else {
+            position_ids_data[i] = static_cast<T>(j++);
+          }
         }
-      }
+      };
+      if (position_ids_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_first_window(position_ids_->GetTensorMutableData<int64_t>());
+      else
+        fill_first_window(position_ids_->GetTensorMutableData<int32_t>());
     }
 
     if (has_mask_input_) {
@@ -509,17 +515,23 @@ void WindowedPositionInputs::Update(DeviceSpan<int32_t> next_tokens, int total_l
       // next_tokens -> [0, a, b, c, d, e]
       // window_size = 3, num_windows = 2, pad_token = 0
       // window_index = 0, attention_mask_ -> ([0] * context_length - window_size_) + [0, 1, 1]
-      auto* attention_mask_data = attention_mask_->GetTensorMutableData<int32_t>();
-      std::fill_n(attention_mask_data, attention_mask_shape_[1] - window_size_, 0);
-      for (size_t i = 0; i < window_size_; i++) {
-        attention_mask_data[attention_mask_shape_[1] - window_size_ + i] = next_tokens.CpuSpan()[i] == model_.config_->model.pad_token_id ? 0 : 1;
-      }
-      for (size_t i = 0; i < window_size_; i++) {
-        if (attention_mask_data[attention_mask_shape_[1] - window_size_ + i] == 1) {
-          attention_mask_backward_offset_ = attention_mask_shape_[1] - window_size_ + i - 1;
-          break;
+      auto fill_first_mask = [&](auto* attention_mask_data) {
+        using T = std::remove_pointer_t<decltype(attention_mask_data)>;
+        std::fill_n(attention_mask_data, attention_mask_shape_[1] - window_size_, T{0});
+        for (size_t i = 0; i < window_size_; i++) {
+          attention_mask_data[attention_mask_shape_[1] - window_size_ + i] = next_tokens.CpuSpan()[i] == model_.config_->model.pad_token_id ? T{0} : T{1};
         }
-      }
+        for (size_t i = 0; i < window_size_; i++) {
+          if (attention_mask_data[attention_mask_shape_[1] - window_size_ + i] == T{1}) {
+            attention_mask_backward_offset_ = attention_mask_shape_[1] - window_size_ + i - 1;
+            break;
+          }
+        }
+      };
+      if (attention_mask_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_first_mask(attention_mask_->GetTensorMutableData<int64_t>());
+      else
+        fill_first_mask(attention_mask_->GetTensorMutableData<int32_t>());
     }
   } else if (window_index_ < num_windows_) {
     if (has_posid_input_) {
@@ -528,9 +540,14 @@ void WindowedPositionInputs::Update(DeviceSpan<int32_t> next_tokens, int total_l
       // window_size = 3, num_windows = 2, pad_token = 0
       // window_index = 1, position_ids_ -> [2, 3, 4]
 
-      auto* position_ids_data = position_ids_->GetTensorMutableData<int32_t>();
-      const auto last_position = position_ids_data[window_size_ - 1];
-      std::iota(position_ids_data, position_ids_data + window_size_, last_position + 1);
+      auto fill_next_window = [&](auto* position_ids_data) {
+        const auto last_position = position_ids_data[window_size_ - 1];
+        std::iota(position_ids_data, position_ids_data + window_size_, last_position + 1);
+      };
+      if (position_ids_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_next_window(position_ids_->GetTensorMutableData<int64_t>());
+      else
+        fill_next_window(position_ids_->GetTensorMutableData<int32_t>());
     }
 
     if (has_mask_input_) {
@@ -538,30 +555,51 @@ void WindowedPositionInputs::Update(DeviceSpan<int32_t> next_tokens, int total_l
       // next_tokens -> [0, a, b, c, d, e]
       // window_size = 3, num_windows = 2, pad_token = 0
       // window_index = 1, attention_mask_ -> ([0] * context_length - (2 * window_size_)) + [0, 1, 1, 1, 1, 1]
-      auto* attention_mask_data = attention_mask_->GetTensorMutableData<int32_t>();
-      std::fill_n(attention_mask_data + attention_mask_backward_offset_ - window_size_ + 1, window_size_, 1);
-      attention_mask_backward_offset_ -= window_size_;
+      auto fill_next_mask = [&](auto* attention_mask_data) {
+        using T = std::remove_pointer_t<decltype(attention_mask_data)>;
+        std::fill_n(attention_mask_data + attention_mask_backward_offset_ - window_size_ + 1, window_size_, T{1});
+        attention_mask_backward_offset_ -= window_size_;
+      };
+      if (attention_mask_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_next_mask(attention_mask_->GetTensorMutableData<int64_t>());
+      else
+        fill_next_mask(attention_mask_->GetTensorMutableData<int32_t>());
     }
   } else {
     // All prompt token chunks have been processed. Now we process the tokens generated by the model.
     if (has_posid_input_) {
       // next_tokens -> [f]
       // position_ids_ -> [5]
-      const auto last_position = position_ids_->GetTensorData<int32_t>()[position_ids_shape_[1] - 1];
-      if (position_ids_shape_[1] != 1) {
-        position_ids_shape_[1] = 1;
-        position_ids_ = OrtValue::CreateTensor(model_.allocator_cpu_, position_ids_shape_, position_ids_type_);
-      }
-      position_ids_->GetTensorMutableData<int32_t>()[0] = last_position + 1;
+      auto fill_generated = [&](auto* data) {
+        using T = std::remove_pointer_t<decltype(data)>;
+        const auto last_position = data[position_ids_shape_[1] - 1];
+        if (position_ids_shape_[1] != 1) {
+          position_ids_shape_[1] = 1;
+          position_ids_ = OrtValue::CreateTensor(model_.allocator_cpu_, position_ids_shape_, position_ids_type_);
+          data = position_ids_->GetTensorMutableData<T>();
+        }
+        data[0] = last_position + 1;
+      };
+      if (position_ids_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_generated(position_ids_->GetTensorMutableData<int64_t>());
+      else
+        fill_generated(position_ids_->GetTensorMutableData<int32_t>());
     }
 
     if (has_mask_input_) {
       // next_tokens -> [f]
       // attention_mask_ -> ([0] * context_length - (2 * window_size_) - 1) + [0, 1, 1, 1, 1, 1, 1]
-      attention_mask_->GetTensorMutableData<int32_t>()[attention_mask_backward_offset_] = 1;
-      if (attention_mask_backward_offset_ > 0) {
-        attention_mask_backward_offset_ -= 1;
-      }
+      auto fill_generated_mask = [&](auto* data) {
+        using T = std::remove_pointer_t<decltype(data)>;
+        data[attention_mask_backward_offset_] = T{1};
+        if (attention_mask_backward_offset_ > 0) {
+          attention_mask_backward_offset_ -= 1;
+        }
+      };
+      if (attention_mask_type_ == Ort::TypeToTensorType<int64_t>)
+        fill_generated_mask(attention_mask_->GetTensorMutableData<int64_t>());
+      else
+        fill_generated_mask(attention_mask_->GetTensorMutableData<int32_t>());
     }
   }
 
diff --git a/src/models/processor.cpp b/src/models/processor.cpp
index 8892f97790..0b64e8b049 100644
--- a/src/models/processor.cpp
+++ b/src/models/processor.cpp
@@ -199,7 +199,20 @@ std::unique_ptr<OrtValue> ProcessTensor<float, int64_t>(OrtxTensor* tensor, Ort:
 }
 
 template std::unique_ptr<OrtValue> ProcessTensor<float>(OrtxTensor* tensor, Ort::Allocator& allocator);
+template std::unique_ptr<OrtValue> ProcessTensor<int32_t>(OrtxTensor* tensor, Ort::Allocator& allocator);
 template std::unique_ptr<OrtValue> ProcessTensor<int64_t>(OrtxTensor* tensor, Ort::Allocator& allocator);
 template std::unique_ptr<OrtValue> ProcessTensor<bool>(OrtxTensor* tensor, Ort::Allocator& allocator);
 
+void EmplaceProcessedTensor(NamedTensors& tensors, std::string_view name,
+                            OrtxTensor* tensor, ONNXTensorElementDataType type,
+                            Ort::Allocator& allocator) {
+  if (type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
+    tensors.emplace(std::string(name), std::make_shared<Tensor>(ProcessTensor<float>(tensor, allocator)));
+  } else if (type == ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16) {
+    tensors.emplace(std::string(name), std::make_shared<Tensor>(ProcessTensor<Ort::BFloat16_t>(tensor, allocator)));
+  } else {
+    tensors.emplace(std::string(name), std::make_shared<Tensor>(ProcessTensor<Ort::Float16_t>(tensor, allocator)));
+  }
+}
+
 }  // namespace Generators
\ No newline at end of file
diff --git a/src/models/processor.h b/src/models/processor.h
index 74a57bb85f..b7c92da2a6 100644
--- a/src/models/processor.h
+++ b/src/models/processor.h
@@ -53,6 +53,12 @@ std::unique_ptr<OrtValue> ProcessTensor(OrtxTensor* tensor, Ort::Allocator& allo
 
 template <typename SrcT, typename DstT>
 std::unique_ptr<OrtValue> ProcessTensor(OrtxTensor* tensor, Ort::Allocator& allocator);
+
+// Helper to emplace a processed tensor with correct type dispatch (float, bf16, fp16)
+void EmplaceProcessedTensor(NamedTensors& tensors, std::string_view name,
+                            OrtxTensor* tensor, ONNXTensorElementDataType type,
+                            Ort::Allocator& allocator);
+
 struct Processor {
   Processor() = default;
   Processor(const Processor&) = delete;