# Gemma4 VLM Investigation — oMLX 0.4.1

*Date: 2026-06-04 | Platform: Apple Silicon M-series, 48 GB, oMLX 0.4.1, mlx-vlm 0.5.0*

---

## Summary

Getting `gemma-4-26b-a4b-it-4bit` to load as a VLM (vision-language model) under oMLX 0.4.1 required fixing **three independent bugs** — none of which produced a clear error message. This document traces the root cause of each, the fix applied, and the validation test.

**End result:** Gemma4 loads as `VLMBatchedEngine`, vision ✅ tool calling ✅ thinking ✅.

---

## Background

The goal was to run a single model that passes all three capability tests required by the `ui-designer` opencode subagent:

1. **Vision** — model must describe image content from a base64 `image_url`
2. **Tool calling** — `finish_reason: tool_calls` with valid JSON args
3. **Thinking** — extended reasoning delivered in `content` or `reasoning_content`

Initial model: `Qwen3` multi-model, multi-modal stack. Abandoned due to unfixed oMLX stream bug (see §A).  
Pivot model: `gemma-4-26b-a4b-it-4bit`. AIO Multimodal; required three config/code fixes (§1–3).

---

## §A — Why Qwen3 Stack Was Abandoned (Background)

**Error:** `"There is no Stream(gpu, 1) in current thread"` on every VLM prefill.

**Root cause:** `_do_external_prefill` in `omlx/scheduler.py` runs on `_mlx_executor` thread without a `with mx.stream(self._stream):` context guard. VLM KV cache states end up referencing a stream from a different thread executor. `mx.eval([c.state])` at line ~2089 fails.

**Status:** Filed as GitHub issue #1636 (opened 2026-06-03). Not fixed in mlx-vlm as of 2026-06-04. Gemma4's stream bug was patched in the upstream PR chain (#1304 → #1485 → #1570); Qwen3-VL was not included.

**Action:** Moved Qwen3 models to `~/.llm/models-disabled/`. Pivoted to Gemma4.

---

## Bug 1 — oMLX Doesn't Detect Gemma4 as VLM

### Symptom

```
omlx.model_discovery - INFO - Model type 'gemma4' is in VLM_MODEL_TYPES but no
vision_config / vit_config / mm_vision_tower found — treating as LLM (text-only quant)
```

Model loads as `BatchedEngine` (text-only) even though it has 358 vision tensor keys.

### Root Cause

`model_discovery.py:_has_vision_subconfig()` checks for exactly three keys:

```python
return (
    "vision_config" in config
    or "vit_config" in config
    or bool(config.get("mm_vision_tower"))
)
```

The mlx-community `gemma-4-26b-a4b-it-4bit` quantized copy had `vision_soft_tokens_per_image: 280` instead of a `vision_config` block (common with re-exported or older quantizations). None of the three checked keys were present.

### Fix

**File:** `/opt/homebrew/Cellar/omlx/0.4.1/libexec/lib/python3.11/site-packages/omlx/model_discovery.py`

```python
# Before
return (
    "vision_config" in config
    or "vit_config" in config
    or bool(config.get("mm_vision_tower"))
)

# After
return (
    "vision_config" in config
    or "vit_config" in config
    or bool(config.get("mm_vision_tower"))
    or "vision_soft_tokens_per_image" in config  # Gemma4 VLM convention
)
```

> **Note:** The longer-term fix is to add `vision_config` to the model's `config.json` (see Bug 2), which makes this patch unnecessary. But it's a useful safety net for any Gemma4 copy missing the block.

### Validation

```
omlx.model_discovery - INFO - Discovered model: gemma-4-26b-a4b-it-4bit (type: vlm, engine: vlm, size: 15.26GB)
```

---

## Bug 2 — VLMBatchedEngine Falls Back to BatchedEngine ("Missing 145 / 752 parameters")

### Symptom

Even after Bug 1 fix, the engine load logged:

```
omlx.engine_pool - WARNING - VLM loading failed for gemma-4-26b-a4b-it-4bit,
  falling back to LLM: Received 145 parameters not in model:
  vision_tower.encoder.layers.16.input_layernorm.weight, ...
omlx.engine_pool - INFO - Successfully loaded gemma-4-26b-a4b-it-4bit as LLM (fallback from VLM)
```

### Root Cause — Part A: Missing `vision_config`

The model's `config.json` had no `vision_config` block. `mlx_vlm/utils.py` does:

```python
config.setdefault("vision_config", {})   # inserts empty dict if absent
config.setdefault("audio_config", {})    # always inserts
```

An empty `vision_config: {}` causes `ModelConfig.from_dict()` to instantiate `VisionConfig()` with **all defaults**:

| Parameter | Default | Actual (26B model) |
|-----------|---------|-------------------|
| `hidden_size` | 768 | **1152** |
| `num_hidden_layers` | 16 | **27** |
| `num_attention_heads` | 12 | **16** |
| `head_dim` | 64 | **72** |
| `intermediate_size` | 3072 | **4304** |
| `standardize` | False | **True** |

With 16-layer default vs 27-layer actual model, weight keys for layers 16–26 have no matching attribute → "Received 145 parameters not in model".

### Root Cause — Part B: `audio_config: {}` Creates AudioEncoder

`config.setdefault("audio_config", {})` always fires, creating a default `AudioConfig()` → `AudioEncoder` with 752 parameters. The model has no audio weights → "Missing 752 parameters: audio_tower.*".

### Fix — Reconstruct `vision_config` in config.json

Parameter values were derived by inspecting the weight tensor shapes:

```bash
# Verification commands used:
python3 -c "
import mlx.core as mx
w = mx.load('model-00003-of-00003.safetensors')
print('hidden_size:', w['vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight'].shape)  # (1152, 1152)
print('head_dim:', w['vision_tower.encoder.layers.0.self_attn.q_norm.weight'].shape)           # (72,)
print('intermediate:', w['vision_tower.encoder.layers.0.mlp.gate_proj.linear.weight'].shape)    # (4304, 1152)
print('has std_bias:', 'vision_tower.std_bias' in w)                                           # True → standardize=True
# Count layers:
layers = set(int(k.split('.')[3]) for k in w if 'vision_tower.encoder.layers.' in k)
print('num_layers:', len(layers), 'max:', max(layers))  # 27, 26 (0-indexed)
"
```

**Add to `~/.llm/models/gemma-4-26b-a4b-it-4bit/config.json`:**

```json
"vision_config": {
  "model_type": "gemma4_vision",
  "hidden_size": 1152,
  "intermediate_size": 4304,
  "num_hidden_layers": 27,
  "num_attention_heads": 16,
  "num_key_value_heads": 16,
  "head_dim": 72,
  "global_head_dim": 72,
  "hidden_activation": "gelu_pytorch_tanh",
  "patch_size": 16,
  "pooling_kernel_size": 3,
  "position_embedding_size": 10240,
  "default_output_length": 280,
  "max_position_embeddings": 131072,
  "rms_norm_eps": 1e-06,
  "rope_parameters": {"rope_theta": 100.0, "rope_type": "default"},
  "standardize": true,
  "use_clipped_linears": false,
  "attention_bias": false,
  "attention_dropout": 0.0
},
"audio_config": null
```

> **`standardize: true` is critical** — without it, vision encoder outputs are un-standardized and the model cannot interpret images. The `std_bias` and `std_scale` weights in the safetensors confirm this flag must be true.
>
> **`audio_config: null`** — `setdefault("audio_config", {})` only fires when the key is *absent*. Setting it explicitly to `null` prevents the empty-dict insertion. Without this, AudioEncoder is created but has no weights → 752 missing parameter warnings and a slower load.

### Validation

```
omlx.engine.vlm - INFO - VLMBatchedEngine loaded: /Users/lvona/.llm/models/gemma-4-26b-a4b-it-4bit
omlx.engine_pool - INFO - Loaded model: gemma-4-26b-a4b-it-4bit (actual: 15.17GB, estimated: 15.26GB)
```

No "fallback from VLM" message. No "Missing parameters" warning.

---

## Bug 3 — Vision Still Not Working (Broken Test Image)

### Symptom

After Bugs 1 & 2 were fixed, model loaded correctly as VLMBatchedEngine but vision test still returned "Please provide the image."

### Root Cause

The test PNG was malformed. oMLX logged:

```
omlx.utils.image - WARNING - Failed to load image: broken PNG file (chunk b'END\xae')
```

The 2×2 pixel base64 string used for testing was constructed incorrectly — a truncated/invalid PNG. `extract_images_from_messages` caught the error, returned zero images, and the model was called with no pixel_values.

### Fix

Use a properly constructed PNG. Python's `struct` + `zlib` produces valid PNGs of any size:

```python
import struct, zlib, base64

def make_png(w, h, rgb):
    def chunk(n, d):
        c = n + d
        return struct.pack('>I', len(d)) + c + struct.pack('>I', zlib.crc32(c) & 0xffffffff)
    raw = b''.join(b'\x00' + bytes(rgb) * w for _ in range(h))
    return (b'\x89PNG\r\n\x1a\n'
            + chunk(b'IHDR', struct.pack('>IIBBBBB', w, h, 8, 2, 0, 0, 0))
            + chunk(b'IDAT', zlib.compress(raw))
            + chunk(b'IEND', b''))

img_b64 = base64.b64encode(make_png(4, 4, [255, 0, 0])).decode()  # red 4×4
```

---

## Final Validation

All three capability tests passing against `gemma-4-26b-a4b-it-4bit` via oMLX :8080:

```
Vision  (3.9s): 'Blue'       → ✅  (correct color identification from 4×4 blue PNG)
Tools   (1.4s): get_weather  → ✅  (finish_reason=tool_calls, args={"location":"Tokyo"})
Thinking(0.4s): '56'         → ✅  (7*8 answered correctly, reasoning in content)
```

---

## Files Changed

### oMLX Source (Homebrew venv — persists until `brew upgrade omlx`)

> ⚠ **Version pin required.** These patches are specific to oMLX **0.4.1**. `setup.sh` calls
> `brew pin omlx` to block upgrades. `llm-start.sh` checks `OMLX_VERSION` in `llm.conf` and
> refuses to start on version mismatch. If you intentionally upgrade, re-run `setup.sh` and
> re-validate all three capability tests (vision, tool calling, thinking) before using.

| File | Change |
|------|--------|
| `omlx/model_discovery.py` | Added `vision_soft_tokens_per_image` to `_has_vision_subconfig()` |
| `omlx/engine_core.py` | Stream/executor creation order fix (for Qwen3 stream bug — prerequisite work, doesn't affect Gemma4) |

### Model Config

| File | Change |
|------|--------|
| `~/.llm/models/gemma-4-26b-a4b-it-4bit/config.json` | Added complete `vision_config` block + `"audio_config": null` |

### opencode Agent Config

| File | Change |
|------|--------|
| `~/.config/opencode/agents/ui-designer.md` | Changed `model:` from `omlx/Qwen3-VL-4B-Instruct-4bit` to `omlx/gemma-4-26b-a4b-it-4bit` |

### oMLX Settings

| File | Change |
|------|--------|
| `~/.omlx/model_settings.json` | Added `model_type_override: "vlm"` for `gemma-4-26b-a4b-it-4bit` (belt-and-suspenders — model_discovery fix makes this redundant but harmless) |

### Models Disabled

| Model | Reason |
|-------|--------|
| `Qwen3-VL-4B-Instruct-4bit` | oMLX stream bug #1636, no upstream fix available |
| `Qwen2-VL-2B-Instruct-4bit` | Same stream bug risk, untested |

---

## Community Research

Sourced from `Blaizzy/mlx-vlm` GitHub issues and PRs (2026-06-04):

- **`vision_soft_tokens_per_image`** is the canonical Gemma4 VLM indicator — official config at `mlx-community/gemma-4-26b-a4b-it-4bit` does include `vision_config`, but locally re-exported copies often strip it.
- **`audio_config: null`** — explicitly null in Google's official config. The 26B and 31B Gemma4 models do not have audio towers. Only the E2B and E4B dense variants support audio.
- **`standardize: true`** — not in the default `VisionConfig` dataclass. Without it, vision outputs are wrong even when the model loads successfully.
- **PR #1027** confirms `gemma-4-26b` continuous batching works; **PR #1261** confirms MTP rollback fix (both merged 2026-06-04 in mlx-vlm 0.6.1 — oMLX 0.4.1 bundles 0.5.0).
- **Issue #1254** warns: always use `apply_chat_template()` for Gemma4 — raw string prompts cause repetition loops. oMLX handles this internally.

---

## Reference: Correct `vision_config` for Gemma4 26B-A4B

Cross-verified against official Google config and mlx-community upload. Use this as ground truth if re-applying after a model re-download:

```json
{
  "model_type": "gemma4_vision",
  "hidden_size": 1152,
  "intermediate_size": 4304,
  "num_hidden_layers": 27,
  "num_attention_heads": 16,
  "num_key_value_heads": 16,
  "head_dim": 72,
  "global_head_dim": 72,
  "hidden_activation": "gelu_pytorch_tanh",
  "patch_size": 16,
  "pooling_kernel_size": 3,
  "position_embedding_size": 10240,
  "default_output_length": 280,
  "max_position_embeddings": 131072,
  "rms_norm_eps": 1e-06,
  "rope_parameters": {"rope_theta": 100.0, "rope_type": "default"},
  "standardize": true,
  "use_clipped_linears": false,
  "attention_bias": false,
  "attention_dropout": 0.0
}
```

---

## §D — Performance Benchmarking and Optimisation

### Throughput baseline (powermode=0, TurboQuant KV 4-bit, M4 Pro 48 GB)

All measurements use streaming with TTFT separated from generation speed (correct methodology). `total_time / output_tokens` is misleading when output is small — prefill dominates and skews the number low.

> **macOS Low Power Mode** throttles GPU and memory controller clocks by ~50%. All tests below used `powermode=0` (normal). Low Power Mode caused 96K and 128K context sizes to fall below the 20 tok/s floor in earlier runs.

| Test | Context in | TTFT | Gen tok/s |
|------|-----------|------|-----------|
| Short text | ~28 tok | 0.4 s | **65.5** |
| Medium text | ~3K tok | 2.4 s | **51.8** |
| Medium text | ~11K tok | 8.6 s | **55.8** |
| Tool calling | ~200 tok | 1.0 s | ✅ `finish=tool_calls` |
| Tool response | ~400 tok | 0.2 s | **58.8** |
| Vision (64×64 PNG) | ~300 tok | 1.2 s | **52.3** |
| 32K context | 32K tok | 30.6 s | **47.4** |
| 64K context | 64K tok | 85.2 s | **41.2** |
| 96K context | 96K tok | 161 s | **35.2** |
| 128K context | 128K tok | 276 s | **31.2** |

All 10/10 pass the ≥20 tok/s generation floor. Model footprint: 15.3 GB weights pinned + 30 GB hard ceiling.

TTFT at large context scales well because `specprefill_enabled: true` uses sparse MoE prefill (activates at 8K+ tokens). In opencode sessions, the prefix cache (SSD) reuses computed prefill states — subsequent turns only pay for new tokens.

### TurboQuant KV vs VLM MTP — head-to-head

Two mutually exclusive optimisation paths were evaluated:

**Path A — TurboQuant KV 4-bit** (`turboquant_kv_enabled: true`):
- 4-bit KV cache compression; ~4× memory reduction vs fp16
- Required to fit 96K+ context within a 30 GB ceiling
- No extra model to load

**Path B — VLM MTP** (`vlm_mtp_enabled: true`) with `mlx-community/gemma-4-26B-A4B-it-assistant-bf16` drafter (801 MB):
- Speculative decoding: drafter proposes tokens, target verifies
- MTP's documented 3.94× gain only applies to batched/concurrent workloads
- Single-user sequential workload shows no speedup; adds 0.9 GB memory pressure
- Cannot coexist with `turboquant_kv_enabled` — hard conflict in `ModelSettings.__post_init__`

**Decision:** TurboQuant KV 4-bit persisted. MTP drafter retained at `~/.llm/drafters/gemma-4-26B-A4B-it-assistant-bf16` for future batched workloads.

### Undocumented oMLX fields discovered (source audit)

From reading oMLX 0.4.1 source (`model_settings.py`, `settings.py`, `scheduler.py`):

**model_settings.json — notable undocumented fields:**
- `min_p` — min-p sampling threshold
- `presence_penalty` — presence penalty
- `max_tool_result_tokens` — cap tool result length
- `forced_ct_kwargs` — lock chat template keys against API override
- `thinking_budget_tokens` — max tokens for thinking block
- `reasoning_parser` — xgrammar parser name (`"qwen"`, `"harmony"`, `"llama"`)
- `guided_grammar` / `guided_grammar_enabled` — EBNF constrained decoding
- `specprefill_enabled` — sparse MoE prefill; **enabled in production** — activates at 8K+ token prompts, measurably reduces TTFT at large context for Gemma4-A4B
- `dflash_enabled` — block-diffusion speculative decoding with SSD cache
- `vlm_mtp_enabled` / `vlm_mtp_draft_model` — VLM MTP (tested above; conflicts with turboquant_kv)
- `model_alias` — expose model under a shorter API name (e.g. `"gemma4"`); **set in production**
- `ttl_seconds` — auto-unload after N idle seconds
- `index_cache_freq` — DSA/SSM model layer cache frequency

**settings.json — undocumented subsections:**
- `scheduler.chunked_prefill: bool` — interleaved prefill/decode for long contexts; **enabled in production**
- `memory.prefill_safe_zone_ratio: float` — secondary prefill guard as fraction of ceiling (default 0.80); **set to 0.90 in production** to match soft_threshold
- `memory.soft_threshold / hard_threshold` — LRU eviction / abort thresholds (default 0.85 / 0.95); **soft_threshold raised to 0.90 in production** — default 0.85 (25.5 GB) caused continuous `adaptive_prefill_throttle` on opencode 22K-token sessions that peaked at 26.9 GB
- `memory.prefill_min_chunk_tokens: int` — minimum tokens per chunked-prefill chunk (default 32); **set to 512 in production** — reduces memory-check overhead 16× per large prefill
- `cache.hot_cache_only: bool` — force all KV blocks in RAM only (useful ≥40 GB); **false in production** (shared system, allow SSD spill)
- `cache.initial_cache_blocks: int` — pre-allocate N cache blocks at startup (default 256); **set to 1024 in production**
- `cache.ssd_cache_dir` — prefix cache directory on SSD; reuses computed prefill states across requests; **set in production** to `~/.llm/prefix-cache`
- `huggingface.hf_cache_enabled: bool` — when true, oMLX scans `~/.cache/huggingface/hub` as a second model directory; **disable in production** to prevent LRU fragmentation