Fix Gemma4 CUDA EP support and audio encoder accuracy (#172)

## Gemma4 multimodal CUDA support & audio encoder fixes

### Changes

**CUDA EP compatibility**
- Lower ONNX opset 24 → 23 for CUDA/TRT EPs (ORT ≤1.24.x lacks opset-24
kernels). Gated by `ort_lower_opset_for_ep` feature flag (default: on).
- Restructure Gemma4 per-layer embedding from one `[V, L*D]` table
(262144×8960 = 2.35B elements) to L separate `[V, D]` tables (262144×256
= 67M each). Avoids ORT CUDA Gather int32 overflow
([onnxruntime#28107](microsoft/onnxruntime#28107))
without any sharding workaround — each layer does its own small Gather
directly.
- Add generic `ort_shard_large_gathers` flag +
`Embedding.shard_weight_dict()` safety net for other models that may hit
the same int32 limit.

**Audio encoder accuracy**
- Implement `ClippableLinear` component with learned input/output
clamping buffers (HF `Gemma4ClippableLinear`). Fixes audio encoder
divergence: max diff 52.68 → 0.0003.
- Fix off-by-one in `_build_causal_window_mask`: `context_left →
context_left - 1` to match HF behavior.

**Example script (`examples/gemma4_multimodal.py`)**
- Fix audio feature extraction: use `Gemma4AudioFeatureExtractor` (not
`WhisperFeatureExtractor`)
- Add `--compare-hf` support for audio and vision+audio modes
- Align `--ep`/`--device` flags with established CLI pattern

### Testing
- All modes (text, vision, audio, vision+audio) verified on CUDA
- Audio `--compare-hf --device cuda`: max diff ~0.0003 (float32)
- Unit tests: 2329 passed (no regressions)
- Lint clean

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>

Author: 4 people

.github/skills/debugging-vl-pipeline/SKILL.md

@@ -1,10 +1,11 @@
 ---
 name: debugging-vl-pipeline
 description: >
-  How to debug vision-language (VL) model output issues in mobius.
-  Covers the systematic pipeline isolation methodology, common failure modes,
-  stage-by-stage comparison with HuggingFace, and numerical tolerance
-  expectations. Use this skill when ORT GenAI multimodal output is wrong,
+  How to debug vision-language (VL) and multimodal (vision + audio) model
+  output issues in mobius.  Covers the systematic pipeline isolation
+  methodology, common failure modes, stage-by-stage comparison with
+  HuggingFace, numerical tolerance expectations, and CUDA EP-specific
+  issues. Use this skill when ORT GenAI multimodal output is wrong,
   garbled, or doesn't match HuggingFace.
 ---
 
@@ -18,6 +19,8 @@ Use this skill when:
 - ONNX model logits diverge significantly from HuggingFace
 - The model generates text-only descriptions ignoring the image
 - Image features appear correct but decoder output is wrong
+- Audio transcription is garbled or wrong despite correct encoder output
+- CUDA EP crashes or produces different results than CPU
 
 ## Debugging methodology: isolate each stage
 
@@ -433,14 +436,98 @@ spatial_merge_size = getattr(vc, "spatial_merge_size", 2)
 temporal_patch_size = getattr(vc, "temporal_patch_size", 2)
 ```
 
+### 6. ClippableLinear divergence (Gemma4)
+
+**Symptoms:** Vision or audio encoder output has large max diff (> 1.0)
+against HuggingFace, even though weights are loaded correctly.
+
+**Root cause:** Gemma4 uses `Gemma4ClippableLinear` with learned finite
+input/output activation clamping for ALL linear layers in its vision and
+audio encoders. Using plain `Linear` misses the clamping.
+
+**Detection:** Check if the HuggingFace model uses `ClippableLinear`:
+```bash
+grep -n "ClippableLinear" transformers/models/<model>/modeling_<model>.py
+```
+
+**Fix:** Use `ClippableLinear` (from `mobius.components`) for all
+affected linear layers. For vision attention: q/k/v/o_proj. For MLP:
+pass `linear_class=ClippableLinear` to the MLP component.
+
+**Impact:**
+- Audio: max diff 52.68 → 0.0003 after fix
+- Vision: max diff 3.92 → 0.00007 after fix
+
+### 7. Missing audio boundary markers
+
+**Symptoms:** Audio transcription is garbled or completely wrong, even
+though the audio encoder output matches HuggingFace.
+
+**Root cause:** HuggingFace wraps audio placeholder tokens with boundary
+markers in `input_ids`:
+```
+<|audio> (256000) + N × <|audio|> (258881) + <audio|> (258883)
+```
+If boundary markers are missing, the model cannot distinguish audio
+regions from text, producing wrong output.
+
+**Fix:** Add boundary markers to `build_input_ids()` when constructing
+audio inputs, matching HuggingFace's token wrapping.
+
+### 8. CUDA EP: ORT Gather int32 overflow
+
+**Symptoms:** CUDA EP crashes or produces incorrect results for models
+with large embedding tables. CPU EP works correctly.
+
+**Root cause:** ORT CUDA `gather_impl.cu` uses `int32` for element
+offset computation: `input_index = idx * cols + col_offset`. For
+tensors with > 2^31 elements (e.g. Gemma4 per-layer embedding
+[262144, 8960] = 2.35B elements), this overflows.
+
+**ORT bug:** microsoft/onnxruntime#28107
+
+**Workaround:** Split large embeddings into smaller tables via
+`nn.ModuleList` so each individual Gather stays under the int32 limit.
+Use Slice instead of Gather for column-wise indexing on large tensors.
+
+### 9. CUDA EP: opset 24 kernel registration
+
+**Symptoms:** ORT CUDA EP fails to find kernels for standard ops
+(Squeeze, Reshape, etc.) even though they work on CPU.
+
+**Root cause:** ORT ≤1.24.x CUDA/TRT EPs don't register kernels for
+opset 24, even though the op semantics are unchanged from opset 23.
+
+**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by
+default) which lowers the declared opset import from 24 to 23 for
+non-CPU EPs. See `src/mobius/_flags.py` and
+`src/mobius/_testing/ort_inference.py`.
+
+### 10. Wrong audio feature extractor
+
+**Symptoms:** Audio model produces completely wrong output. Audio
+encoder features don't match HuggingFace at all.
+
+**Root cause:** Using `WhisperFeatureExtractor` instead of the
+model-specific feature extractor (e.g. `Gemma4AudioFeatureExtractor`).
+Different extractors produce different mel spectrograms.
+
+**Fix:** Always use the correct feature extractor for the model:
+```python
+from transformers import AutoFeatureExtractor
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
+```
+
 ## Reference files
 
 - **Integration tests:** `tests/integration_test.py`
   (`TestVLFullForward`, `TestQwen25VL3Model`)
 - **ORT GenAI tests:** `tests/ort_genai_test.py`
   (`TestOrtGenaiQwen25VL.test_multimodal_image_generation`)
 - **Example scripts:** `examples/qwen25_vl_ort_genai.py`,
-  `examples/qwen3_vl_ort_genai.py`
+  `examples/qwen3_vl_ort_genai.py`, `examples/gemma4_multimodal.py`
 - **genai_config reference:** `.github/skills/ort-genai-config/SKILL.md`
 - **ORT GenAI position_ids code (external):**
   `onnxruntime-genai/src/models/position_inputs.cpp:617-814`
+- **Feature flags:** `src/mobius/_flags.py`
+  (`ort_lower_opset_for_ep`, `ort_cuda_grouped_rmsnorm_workaround`)

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

@@ -1,20 +1,21 @@
 ---
 name: multimodal-models
 description: >
-  How to add multimodal (vision + language) models to mobius.
+  How to add multimodal (vision + language + audio) models to mobius.
   Covers projector variants (Gemma3, MLP, Linear), the VisionModel encoder,
-  InputMixer, VisionLanguageTask, image token handling, and weight name
-  mappings. Use this skill when adding a model that processes both images and
-  text.
+  InputMixer, VisionLanguageTask, image/audio token handling, ClippableLinear,
+  and weight name mappings. Use this skill when adding a model that processes
+  images, audio, or both alongside text.
 ---
 
-# Skill: Multimodal (Vision + Language) Models
+# Skill: Multimodal (Vision + Language + Audio) Models
 
 ## When to use
 
 Use this skill when adding a model that processes both images and text — such
-as Gemma3, LLaVA, LLaVA-NeXT, Phi-3-Vision, PaliGemma, InternVL2, Pixtral,
-Idefics2/3, Molmo, Florence2, or Video-LLaVA.
+as Gemma3, Gemma4, LLaVA, LLaVA-NeXT, Phi-3-Vision, PaliGemma, InternVL2,
+Pixtral, Idefics2/3, Molmo, Florence2, or Video-LLaVA — or a model that
+also processes audio (e.g. Gemma4 with speech/audio inputs).
 
 ## Architecture overview
 
@@ -390,6 +391,90 @@ The vision pipeline is completely shared with Qwen3-VL — only the text
 decoder differs (hybrid DeltaNet + full attention). This means vision
 encoder bugs/fixes apply to both models equally.
 
+## Gemma4: vision + audio multimodal
+
+Gemma4 models (E2B, E4B, 26B-A4B, 31B) support **both vision and audio**
+inputs. The architecture has 4 sub-models: decoder, vision encoder, audio
+encoder, and embedding.
+
+### Architecture
+
+```
+pixel_values ──► [Vision Encoder] ──► image_features ──┐
+                                                        │
+audio_features ─► [Audio Encoder] ──► audio_features ──┤
+                                                        │
+input_ids ──────► [Embedding] ◄────────────────────────┘
+                       │
+                       ▼
+                 [Text Decoder] ──► logits
+```
+
+### ClippableLinear (critical for Gemma4)
+
+Gemma4's vision and audio encoders use `Gemma4ClippableLinear` — a
+`Linear` with learned finite input/output activation clamping:
+
+```python
+x = Clip(x, input_min, input_max)
+x = x @ weight.T [+ bias]
+x = Clip(x, output_min, output_max)
+```
+
+**This is the single most common source of Gemma4 divergence.** Using
+plain `Linear` instead of `ClippableLinear` causes:
+- Audio encoder max diff: 52.68 → 0.0003 after fix
+- Vision encoder max diff: 3.92 → 0.00007 after fix
+
+Vision encoder uses ClippableLinear for ALL linear layers:
+- Q/K/V/O projections in `Gemma4VisionSelfAttention`
+- gate/up/down projections in MLP (via `linear_class=ClippableLinear`)
+
+Audio encoder uses ClippableLinear for its linear layers as well.
+
+See `reusable-components` skill for full `ClippableLinear` API reference.
+
+### Audio boundary markers
+
+HuggingFace wraps audio tokens with boundary markers that must be present
+in `input_ids` for correct generation:
+
+```
+<|audio>  (256000) + N × <|audio|> (258881) + <audio|> (258883)
+```
+
+This parallels the image token pattern:
+```
+<|image>  (255999) + N × <|image|> (258880) + <image|> (258882)
+```
+
+**Missing audio boundary markers** cause garbled audio transcription output
+even when the audio encoder output is numerically correct.
+
+### Per-layer embeddings (CUDA ORT workaround)
+
+Gemma4 uses per-layer embedding: `embed_tokens_per_layer` with shape
+`[V, L*D]` where V=vocab_size, L=num_layers, D=per_layer_dim. For large
+models this creates a single Gather on a 2.35B-element tensor, which
+**overflows ORT's CUDA Gather kernel** (int32 offset computation in
+`gather_impl.cu`).
+
+**Workaround:** Split into L separate `Embedding([V, D])` tables via
+`nn.ModuleList`. In `preprocess_weights`, split the HF weight column-wise:
+```python
+for i in range(num_layers):
+    renamed[f"embed_tokens_per_layer.{i}.weight"] = value[:, i*D:(i+1)*D]
+```
+
+Use `Slice` instead of `Gather` for per-layer projection indexing to
+avoid the large-tensor issue entirely.
+
+### Audio feature extraction
+
+Gemma4 uses `Gemma4AudioFeatureExtractor` (not `WhisperFeatureExtractor`).
+Using the wrong feature extractor produces completely different mel features
+and the model fails silently (produces garbage transcription).
+
 ## Testing multimodal models
 
 ### Image token count

.github/skills/phi4mm-component-parity/SKILL.md

@@ -5,9 +5,10 @@ description: >
   embedding/projector, text decoder) matches HuggingFace output. Covers pipeline
   isolation methodology, common failure modes from real debugging experience,
   step-by-step debugging process, and integration test patterns. Applicable to
-  any multimodal model with similar architecture (Phi4MM, future audio+vision
-  models). Use this skill when multimodal ONNX model output diverges from
-  HuggingFace, or when adding a new multimodal model with multiple encoders.
+  any multimodal model with similar architecture (Phi4MM, Gemma4, future
+  audio+vision models). Use this skill when multimodal ONNX model output
+  diverges from HuggingFace, or when adding a new multimodal model with
+  multiple encoders.
 ---
 
 # Skill: Multimodal Component Parity Debugging
@@ -371,6 +372,46 @@ the actual sequence includes fused image/audio tokens, the lengths diverge.
 **Fix:** Always derive `seq_len` from `inputs_embeds.shape[1]` when the
 model uses inputs_embeds as input.
 
+### 8. ClippableLinear not used in encoder (Gemma4)
+
+**Symptoms:** Encoder output has large numerical divergence (max diff > 1.0)
+from HuggingFace, despite all weights loading correctly.
+
+**Root cause:** Some HuggingFace models (e.g. Gemma4) use
+`ClippableLinear` — a linear layer with learned finite input/output
+activation clipping — for ALL linear layers in their vision and audio
+encoders (attention q/k/v/o projections AND MLP gate/up/down projections).
+Using plain `Linear` misses the clamping and causes divergence.
+
+**Detection:** Check HF source for `ClippableLinear`:
+```bash
+grep -n "ClippableLinear" transformers/models/<model>/modeling_<model>.py
+```
+
+**Fix:** Use `ClippableLinear` from `mobius.components`:
+- For attention: use `ClippableLinear` for q/k/v/o_proj
+- For MLP: pass `linear_class=ClippableLinear` parameter
+
+**Impact (Gemma4):**
+- Audio: max diff 52.68 → 0.0003
+- Vision: max diff 3.92 → 0.00007
+
+### 9. Missing audio/image boundary tokens
+
+**Symptoms:** Audio transcription or vision description is garbled, but
+encoder output matches HuggingFace.
+
+**Root cause:** HuggingFace wraps modality placeholder tokens with
+boundary markers:
+- Image: `<|image>` (open) + N × `<|image|>` (pad) + `<image|>` (close)
+- Audio: `<|audio>` (open) + N × `<|audio|>` (pad) + `<audio|>` (close)
+
+Missing boundary markers prevent the model from correctly identifying
+modality regions in the input sequence.
+
+**Fix:** Ensure `build_input_ids` wraps placeholder tokens with the
+correct open/close marker token IDs.
+
 ## Step-by-step debugging process
 
 ### Phase 1: Text-only baseline
@@ -562,11 +603,16 @@ image_sizes = np.array([[384, 384]], dtype=np.int64)
 - **Integration tests:** `tests/phi4mm_integration_test.py`,
   `tests/integration_test.py`
 - **VL debugging skill:** `.github/skills/debugging-vl-pipeline/SKILL.md`
-- **Model implementation:** `src/mobius/models/phi.py`
-- **Audio components:** `src/mobius/components/_audio.py`
+- **Model implementation:** `src/mobius/models/phi.py`,
+  `src/mobius/models/gemma4.py`
+- **Audio components:** `src/mobius/components/_audio.py`,
+  `src/mobius/components/_gemma4_audio.py`
 - **Vision components:** `src/mobius/components/_vision.py`
+- **ClippableLinear:** `src/mobius/components/_gemma4_audio.py`
 - **LoRA component:** `src/mobius/components/_lora.py`
 - **Weight loading:** `src/mobius/_weight_loading.py`
+- **Feature flags:** `src/mobius/_flags.py`
 - **ORT GenAI config skill:** `.github/skills/ort-genai-config/SKILL.md`
 - **Weight name alignment skill:**
   `.github/skills/weight-name-alignment/SKILL.md`
+- **Gemma4 example:** `examples/gemma4_multimodal.py`

.github/skills/quality-checklist/SKILL.md

@@ -87,6 +87,14 @@ before the PR is merged.
       `testdata/golden/<cat>/<model>_generation.json`
 - [ ] `python -m pytest tests/e2e_golden_test.py -m generation -k "<model>"` passes
 
+> **Speech-language models:** The golden generation script supports the
+> `speech-language` task type for models that process audio inputs (e.g.,
+> Gemma4). The `_generate_speech_language()` function in
+> `scripts/generate_golden.py` handles audio feature extraction and input
+> construction for these models. Ensure the correct feature extractor
+> (e.g. `Gemma4AudioFeatureExtractor`, not `WhisperFeatureExtractor`) is
+> auto-detected via `AutoFeatureExtractor.from_pretrained()`.
+
 > **Why L4/L5 matter:** Graph-build tests (L1) only verify ONNX graph
 > construction; they never execute the graph with real data.  A MatMul shape
 > mismatch that crashes at runtime, a wrong normalisation type, or a missing
@@ -113,6 +121,15 @@ python examples/<model>_text_generation.py --compare-hf --dtype bf16
 - [ ] `mobius build --model <hf-model-id> /tmp/out` completes without error
 - [ ] Output directory contains the expected ONNX files and `genai_config.json`
 
+### 8a. Multi-EP correctness (CUDA)
+
+- [ ] Model runs correctly with `--ep cuda` (or `--device cuda`)
+- [ ] CUDA results match CPU results (compare generation output)
+- [ ] No crashes from large tensor operations (see ORT Gather int32
+      overflow: microsoft/onnxruntime#28107)
+- [ ] `ort_lower_opset_for_ep` flag handles opset 24→23 lowering for
+      CUDA EP (enabled by default in `src/mobius/_flags.py`)
+
 ### 9. ORT GenAI runtime
 
 - [ ] Model can be loaded with `ort_genai.Model(output_dir)` without error