feat: add Gemma 4 support (text, vision, audio, MoE) (#137)

Adds full Gemma 4 support for all four production variants.

## Model Variants

| Model ID | Type | Modalities |
|---|---|---|
| `google/gemma-4-E2B-it` | Any-to-Any | Vision + Audio + Text |
| `google/gemma-4-E4B-it` | Any-to-Any | Vision + Audio + Text |
| `google/gemma-4-26B-A4B-it` | Image-Text | Vision + Text |
| `google/gemma-4-31B-it` | Image-Text | Vision + Text |

## Architecture

- **Dual RoPE**: sliding-window layers use `rope_type=default, θ=10k`;
full-attention layers use `rope_type=proportional, partial_factor=0.25,
θ=1M`
- **5:1 sliding/full pattern**: every 6th layer is full attention
(`layer_types` list from HF config)
- **GQA**: 8 Q heads / 1 KV head; `head_dim=256` for sliding layers,
`global_head_dim=512` for full-attention layers
- **KV sharing** (`num_kv_shared_layers=20` for E2B/E4B): last N layers
receive K,V as graph edges from source layers — no `k_proj`/`v_proj`
weights allocated
- **Double-wide MLP** (`use_double_wide_mlp=True`): KV-shared layers use
`2x intermediate_size`
- **Final logit softcapping**: `logit_cap * tanh(x / logit_cap)` with
`logit_cap=30.0`
- **Per-layer input embeddings**: optional gating from
`vocab_size_per_layer_input` embeddings
- **`layer_scalar`**: per-layer learned scalar applied after per-layer
input contribution

## Audio (E2B/E4B only)

- `Gemma4AudioEncoder`: Conformer-based encoder with 12 layers,
`hidden_size=1024`
- Causal chunked attention (`attention_chunk_size=12`,
`context_left=13`)
- Convolutional subsampler with channels `[128, 32]` (4x temporal
downsampling)
- Projects to text `hidden_size` via `output_proj_dims=1536`

## MoE (26B-A4B and 31B)

- Direct `com.microsoft.MoE` op emission in `Gemma4DecoderLayer` when
`enable_moe_block=True` and EP supports MoE
- Falls back to dense MLP when EP does not support MoE
- Gate router matches HF `Gemma4SparseMoeBlock` routing

## Tasks

- **`Gemma4VisionLanguageTask`** (3-model split): `decoder` + `vision` +
`embedding` — for 26B-A4B, 31B
- **`Gemma4AnyToAnyTask`** (4-model split): `decoder` + `vision` +
`audio` + `embedding` — for E2B, E4B

## Example

`examples/gemma4_multimodal.py` — end-to-end Any-to-Any generation with
interleaved image/audio/text.

## Tests

- `test_graph_builds_without_weights[gemma4_text]` + 3 parametrized
variants
- `test_gemma4_multimodal_graph` — 3-model split (decoder + vision +
embedding)
- `test_gemma4_any_to_any_graph` — 4-model split with
`num_kv_shared_layers=1`, audio encoder
- **2377 total tests passing**

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

.github/skills/multimodal-models/SKILL.md

@@ -534,45 +534,48 @@ gathered = op.Gather(padded, indices, axis=0)
 result = op.Where(image_mask, gathered, text_embeddings)
 ```
 
-## Any-to-Any 4-model task (vision+audio+text)
+## Conditional 3-or-4-model task (vision+audio+text)
 
 Some models come in two tiers: small variants support vision **and** audio
 (Any-to-Any), while large variants support vision only (Image-Text-to-Text).
-Each tier uses a different number of exported ONNX models.
+A **single unified task class** handles both tiers by checking whether
+`config.audio is not None` to decide whether to include the speech encoder.
 
 ### Tier split
 
 | Tier | Models | ONNX split | Example |
 |------|--------|-----------|---------|
-| Small Any-to-Any | E2B, E4B | 4 models: decoder + vision + **audio** + embedding | `google/gemma-4-E2B-it` |
+| Small Any-to-Any | E2B, E4B | 4 models: decoder + vision + **speech** + embedding | `google/gemma-4-E2B-it` |
 | Large Image-Text-to-Text | 26B-A4B, 31B | 3 models: decoder + vision + embedding | `google/gemma-4-26B-A4B-it` |
 
-The task class detects which tier to use from the config (e.g. whether
-`config.audio is not None` — `ArchitectureConfig.from_transformers` populates
-the `audio` field when the HuggingFace config contains an audio sub-config).
+`ArchitectureConfig.from_transformers` populates the `audio` field when the
+HuggingFace config contains an audio sub-config; otherwise it is `None`.
 
-### 4-model task structure
+### 4-model task structure (when audio is present)
 
 ```
-decoder          inputs_embeds [B, S, H] → logits + KV cache
-vision_encoder   pixel_values [B, 3, H, W] → image_features [num_image_tokens, H]
-audio_encoder    input_features [B, T, mel] → audio_features [num_audio_tokens, H]
-embedding        input_ids + image_features + audio_features → inputs_embeds [B, S, H]
+decoder   inputs_embeds [B, S, H] → logits + KV cache
+vision    pixel_values [B, N, 3*P^2] → image_features [B*N, H]
+speech    input_features [B, T, mel] → audio_features [num_audio_tokens, H]
+embedding input_ids + image_features + audio_features → inputs_embeds [B, S, H]
 ```
 
-Reference implementation: `Gemma4AnyToAnyTask` in
-`src/mobius/tasks/_gemma4.py`. This follows the same 4-model structural pattern as
+Reference implementation: `Gemma4Task` in `src/mobius/tasks/_gemma4.py`.
+This follows the same multi-model structural pattern as
 `Phi4MMMultiModalTask` in `src/mobius/tasks/_phi4mm_multimodal.py`
 (each modality is a separate ONNX model; embedding splices features at placeholder
 positions), though the exact I/O shapes differ per architecture.
 
-### Audio encoder wiring
+`Gemma4VisionLanguageTask` and `Gemma4AnyToAnyTask` are backward-compatible
+aliases pointing to `Gemma4Task`.
 
-The audio encoder takes raw mel-spectrogram frames and outputs token-level
-features at the text hidden size:
+### Speech encoder wiring
+
+The speech (audio) encoder takes raw mel-spectrogram frames and outputs
+token-level features at the text hidden size:
 
 ```python
-# In Gemma4AnyToAnyTask._build_audio():
+# In Gemma4Task._build_speech():
 input_features = ir.Value(
     name="input_features",
     shape=ir.Shape([batch, time, input_size]),   # [B, T, 128]
@@ -583,37 +586,37 @@ audio_features = audio_encoder(op, input_features)
 ```
 
 The audio encoder (`Gemma4AudioEncoder` / `_Gemma4AudioEncoderModel`) is
-its own `nn.Module` subgraph exported as the `"audio"` key in the
+its own `nn.Module` subgraph exported as the `"speech"` key in the
 `ModelPackage`.
 
 ### Embedding model fuses all modalities
 
-The embedding model receives `input_ids`, `image_features`, and
-`audio_features` as separate inputs and splices them into the token
-embedding sequence at the placeholder positions:
+The embedding model receives `input_ids`, `image_features`, and optionally
+`audio_features` as inputs and splices them into the token embedding sequence
+at the placeholder positions:
 
 ```python
-# In Gemma4AnyToAnyTask._build_embedding():
+# In Gemma4Task._build_embedding():
 inputs_embeds = embedding(
     op,
     input_ids=input_ids,            # [B, S]
     image_features=image_features,  # [num_image_tokens, H]
-    audio_features=audio_features,  # [num_audio_tokens, H]
+    audio_features=audio_features,  # [num_audio_tokens, H] — only when audio present
 )
 # returns inputs_embeds: [B, S, H]
 ```
 
-### Task class tier detection
+### Task class conditional pattern
 
 ```python
-class MyAnyToAnyTask(ModelTask):
+class MyConditionalTask(ModelTask):
     def build(self, module, config):
         models = {}
         models["decoder"] = self._build_decoder(module.decoder, config)
         models["vision"] = self._build_vision(module.vision_encoder, config)
-        models["embedding"] = self._build_embedding(module.embedding, config)
-        # Build audio encoder only when audio config is present
+        # Build speech encoder only when audio config is present
         if config.audio is not None:
-            models["audio"] = self._build_audio(module.audio_encoder, config)
+            models["speech"] = self._build_speech(module.audio_encoder, config)
+        models["embedding"] = self._build_embedding(module.embedding, config)
         return ModelPackage(models, config=config)
 ```

examples/gemma4/ort_genai/README.md

@@ -0,0 +1,182 @@
+# Gemma 4 ORT GenAI Config Files
+
+ORT GenAI configuration files for the **`google/gemma-4-E2B-it`** checkpoint
+(Gemma 4 Embedding-Enhanced 2B instruction-tuned model).
+
+> **Important:** The `gemma4`, `gemma4_text`, and `gemma4_any_to_any` model
+> types are not yet in a released ORT GenAI build. These configs require ORT
+> GenAI support for the Gemma 4 architecture. See the *ORT GenAI Support*
+> section below.
+
+---
+
+## Directory layout
+
+```
+ort_genai/
+├── text/
+│   └── genai_config.json        # Text-only decoder (model.onnx)
+├── vlm/
+│   ├── genai_config.json        # 3-model VLM (model.onnx + vision.onnx + embedding.onnx)
+│   └── processor_config.json    # SigLIP image processor config
+└── any_to_any/
+    └── genai_config.json        # 4-model AnyToAny (+speech.onnx)
+```
+
+---
+
+## Key architecture values (`google/gemma-4-E2B-it`)
+
+| Field | Value |
+|---|---|
+| `vocab_size` | 262 144 |
+| `hidden_size` | 1 536 |
+| `num_attention_heads` | 8 |
+| `num_key_value_heads` | 1 |
+| `head_dim` (local/sliding layers) | 256 |
+| `global_head_dim` (full-attention layers) | 512 |
+| `num_hidden_layers` (total) | 35 |
+| `num_kv_shared_layers` | 20 |
+| **KV cache depth** | **15** (= 35 − 20) |
+| `sliding_window` | 512 |
+| `max_position_embeddings` | 131 072 |
+| `bos_token_id` | 2 |
+| `eos_token_id` | `[1, 106]` |
+| `image_token_id` | 255 999 (`boi_token_id`) |
+| `audio_token_id` | 258 881 |
+| Vision patch size | 16 |
+| Vision tokens per image | 280 |
+
+---
+
+## KV cache sharing (15, not 35, layers)
+
+Gemma 4 uses **KV projection sharing**: the last `num_kv_shared_layers = 20`
+decoder layers reuse the K and V projections from the preceding layers.
+Those 20 layers therefore have **no independent KV cache entries**.
+
+Only the first `35 − 20 = 15` layers produce their own KV cache, so:
+
+- `num_hidden_layers` in `genai_config.json` is **15**, not 35.
+- `past_key_values.{0..14}.key/value` are the only KV inputs/outputs.
+
+ORT GenAI must be aware of this sharing pattern to feed the correct cached
+K/V values to the shared layers during autoregressive decoding.
+
+---
+
+## Sliding window attention (5:1 local:global pattern)
+
+The 35-layer stack alternates 4 sliding-window (local) layers followed by
+1 full-attention (global) layer, repeated 7 times:
+
+```
+layers 0–3   : sliding_attention  (window = 512 tokens)
+layer  4     : full_attention
+layers 5–8   : sliding_attention
+layer  9     : full_attention
+layers 10–13 : sliding_attention
+layer 14     : full_attention
+(layers 15–34 share KV — no independent cache)
+```
+
+The `sliding_window.layers` field in the configs lists the **sliding** layer
+indices within the 15-entry KV cache: `[0, 1, 2, 3, 5, 6, 7, 8, 10, 11, 12, 13]`.
+
+### Dual head_dim
+
+Sliding (local) attention layers use `head_dim = 256`; full (global)
+attention layers use `global_head_dim = 512`. The `head_size` field in
+`genai_config.json` is set to **256** (the local value). ORT GenAI support
+for the Gemma 4 model type must handle the per-layer head_dim difference
+internally.
+
+---
+
+## Vision encoding
+
+Gemma 4 uses a **SigLIP** vision encoder (ViT, patch_size=16, image_size=448)
+with a "pan-and-scan" tiling strategy for high-resolution images.
+
+The ONNX vision model (`vision.onnx`) takes **pre-patchified** inputs:
+- `pixel_values [batch, num_patches, 3 * 16 * 16]` — flattened patch pixels
+- `pixel_position_ids [batch, num_patches, 2]` — (row, col) patch coordinates
+
+This differs from other VLMs that pass raw images and use `image_grid_thw`.
+The HuggingFace `AutoProcessor` produces the pre-patchified format directly.
+
+Each image produces **280 soft tokens** (`vision_soft_tokens_per_image = 280`).
+
+---
+
+## Audio encoding
+
+Gemma 4 Any-to-Any uses a **Conformer** audio encoder (12 layers,
+hidden_size=1024) with 4× subsampling.
+
+The ONNX audio model (`speech.onnx`) takes:
+- `input_features [batch, time, 128]` — 128-dim mel spectrogram
+
+Output: `audio_features [batch, time/4, 1536]` — projected to text hidden_size.
+
+`audio_token_id = 258881` identifies audio soft-token positions in `input_ids`.
+
+---
+
+## ORT GenAI support required
+
+These configs use model types not yet in a released ORT GenAI build:
+
+| Config | `model.type` | Required new type |
+|---|---|---|
+| `text/genai_config.json` | `gemma4_text` | New decoder-only type |
+| `vlm/genai_config.json` | `gemma4` | New VLM type |
+| `any_to_any/genai_config.json` | `gemma4_any_to_any` | New multimodal type |
+
+Until these types are supported, you can approximate text-only generation
+using `"type": "gemma3_text"` (without KV sharing or dual head_dim support).
+
+---
+
+## Usage example
+
+```python
+import onnxruntime_genai as og
+
+# Text-only
+model = og.Model("path/to/gemma4/ort_genai/text")
+tokenizer = og.Tokenizer(model)
+params = og.GeneratorParams(model)
+params.set_search_options(max_length=512)
+generator = og.Generator(model, params)
+
+# VLM (requires vision.onnx + embedding.onnx alongside model.onnx)
+model = og.Model("path/to/gemma4/ort_genai/vlm")
+processor = og.MultiModalProcessor(model)
+# ... load image and tokenize prompt with processor ...
+```
+
+The model directory must contain the ONNX files exported from mobius:
+```
+text/
+  model.onnx          ← exported by Gemma4TextCausalLMTask
+  tokenizer.json
+  genai_config.json
+
+vlm/
+  model.onnx          ← decoder (Gemma4VisionLanguageTask)
+  vision.onnx         ← vision encoder
+  embedding.onnx      ← embedding fusion
+  tokenizer.json
+  genai_config.json
+  processor_config.json
+
+any_to_any/
+  model.onnx          ← decoder (Gemma4AnyToAnyTask)
+  vision.onnx
+  speech.onnx         ← Conformer audio encoder
+  embedding.onnx
+  tokenizer.json
+  genai_config.json
+  processor_config.json  ← copy from vlm/
+```

examples/gemma4/ort_genai/any_to_any/genai_config.json

@@ -0,0 +1,100 @@
+{
+    "model": {
+        "type": "gemma4_any_to_any",
+        "vocab_size": 262144,
+        "context_length": 131072,
+        "bos_token_id": 2,
+        "eos_token_id": [1, 106],
+        "pad_token_id": 0,
+        "image_token_id": 255999,
+        "audio_token_id": 258881,
+        "decoder": {
+            "session_options": {
+                "log_id": "onnxruntime-genai",
+                "provider_options": []
+            },
+            "filename": "model.onnx",
+            "hidden_size": 1536,
+            "head_size": 256,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 1,
+            "num_hidden_layers": 15,
+            "inputs": {
+                "inputs_embeds": "inputs_embeds",
+                "attention_mask": "attention_mask",
+                "position_ids": "position_ids",
+                "past_key_names": "past_key_values.%d.key",
+                "past_value_names": "past_key_values.%d.value"
+            },
+            "outputs": {
+                "logits": "logits",
+                "present_key_names": "present.%d.key",
+                "present_value_names": "present.%d.value"
+            },
+            "sliding_window": {
+                "window_size": 512,
+                "pad_value": 0,
+                "alignment": "right",
+                "slide_key_value_cache": true,
+                "slide_inputs": true,
+                "layers": [0, 1, 2, 3, 5, 6, 7, 8, 10, 11, 12, 13]
+            }
+        },
+        "vision": {
+            "filename": "vision.onnx",
+            "config_filename": "processor_config.json",
+            "session_options": {
+                "log_id": "onnxruntime-genai",
+                "provider_options": []
+            },
+            "inputs": {
+                "pixel_values": "pixel_values",
+                "pixel_position_ids": "pixel_position_ids"
+            },
+            "outputs": {
+                "image_features": "image_features"
+            }
+        },
+        "speech": {
+            "filename": "speech.onnx",
+            "session_options": {
+                "log_id": "onnxruntime-genai",
+                "provider_options": []
+            },
+            "inputs": {
+                "input_features": "input_features"
+            },
+            "outputs": {
+                "audio_features": "audio_features"
+            }
+        },
+        "embedding": {
+            "filename": "embedding.onnx",
+            "session_options": {
+                "log_id": "onnxruntime-genai",
+                "provider_options": []
+            },
+            "inputs": {
+                "input_ids": "input_ids",
+                "image_features": "image_features",
+                "audio_features": "audio_features"
+            },
+            "outputs": {
+                "inputs_embeds": "inputs_embeds"
+            }
+        }
+    },
+    "search": {
+        "do_sample": false,
+        "early_stopping": true,
+        "max_length": 8192,
+        "min_length": 0,
+        "num_beams": 1,
+        "num_return_sequences": 1,
+        "past_present_share_buffer": false,
+        "repetition_penalty": 1.0,
+        "temperature": 1.0,
+        "top_k": 1,
+        "top_p": 1.0
+    }
+}