[object Object]

← back to Exo

Add Gemma 4 + VLM fixes + thinking parsing updates (#1851)

196543ce6922a023ac240b7f43031e8d621ffdeb · 2026-04-11 12:29:33 +0100 · rltakashige

## Motivation
Add support for Gemma 4, including VLM!

## Changes

- Add auto parallel strategies and model cards for Gemma 4
- Normalise Gemma 4's special Vision Transformer handling to be in line
with the rest of our vision processors.
- Also adds reprs to messages and b64 hashes to prevent log spam.

## Test Plan

### Manual Testing
Tested manually on 4bit E2B and 8bit 26B

### Automated Testing
Model onboarding shows small logit diffs.

---------

Co-authored-by: Evan <evanev7@gmail.com>

Files touched

Diff

commit 196543ce6922a023ac240b7f43031e8d621ffdeb
Author: rltakashige <rl.takashige@gmail.com>
Date:   Sat Apr 11 12:29:33 2026 +0100

    Add Gemma 4 + VLM fixes + thinking parsing updates (#1851)
    
    ## Motivation
    Add support for Gemma 4, including VLM!
    
    ## Changes
    
    - Add auto parallel strategies and model cards for Gemma 4
    - Normalise Gemma 4's special Vision Transformer handling to be in line
    with the rest of our vision processors.
    - Also adds reprs to messages and b64 hashes to prevent log spam.
    
    ## Test Plan
    
    ### Manual Testing
    Tested manually on 4bit E2B and 8bit 26B
    
    ### Automated Testing
    Model onboarding shows small logit diffs.
    
    ---------
    
    Co-authored-by: Evan <evanev7@gmail.com>
---
 .mlx_typings/mlx/nn/layers/quantized.pyi           |   8 +-
 .mlx_typings/mlx_lm/models/gemma4.pyi              |  31 +++
 .mlx_typings/mlx_lm/models/gemma4_text.pyi         | 179 +++++++++++++++
 .mlx_typings/mlx_lm/tokenizer_utils.pyi            |   2 +
 .../mlx-community--gemma-4-26b-a4b-it-4bit.toml    |  15 ++
 .../mlx-community--gemma-4-26b-a4b-it-6bit.toml    |  15 ++
 .../mlx-community--gemma-4-26b-a4b-it-8bit.toml    |  15 ++
 .../mlx-community--gemma-4-26b-a4b-it-bf16.toml    |  15 ++
 .../mlx-community--gemma-4-31b-it-4bit.toml        |  15 ++
 .../mlx-community--gemma-4-31b-it-6bit.toml        |  15 ++
 .../mlx-community--gemma-4-31b-it-8bit.toml        |  15 ++
 .../mlx-community--gemma-4-31b-it-bf16.toml        |  15 ++
 .../mlx-community--gemma-4-e2b-it-4bit.toml        |  15 ++
 .../mlx-community--gemma-4-e2b-it-6bit.toml        |  15 ++
 .../mlx-community--gemma-4-e2b-it-8bit.toml        |  15 ++
 .../mlx-community--gemma-4-e2b-it-bf16.toml        |  15 ++
 .../mlx-community--gemma-4-e4b-it-4bit.toml        |  15 ++
 .../mlx-community--gemma-4-e4b-it-6bit.toml        |  15 ++
 .../mlx-community--gemma-4-e4b-it-8bit.toml        |  15 ++
 .../mlx-community--gemma-4-e4b-it-bf16.toml        |  15 ++
 src/exo/shared/models/model_cards.py               |   3 +-
 src/exo/worker/engines/mlx/auto_parallel.py        |   2 +-
 src/exo/worker/engines/mlx/cache.py                |   2 +-
 .../worker/engines/mlx/generator/batch_generate.py |   6 +-
 src/exo/worker/engines/mlx/generator/generate.py   |  61 ++++-
 src/exo/worker/engines/mlx/utils_mlx.py            |  52 +++--
 src/exo/worker/engines/mlx/vision.py               | 255 ++++++++++++++++++---
 src/exo/worker/main.py                             |   4 +-
 .../runner/llm_inference/model_output_parsers.py   |  18 +-
 uv.lock                                            |   6 +-
 30 files changed, 795 insertions(+), 74 deletions(-)

diff --git a/.mlx_typings/mlx/nn/layers/quantized.pyi b/.mlx_typings/mlx/nn/layers/quantized.pyi
index 137a4c8e..7cd43dd6 100644
--- a/.mlx_typings/mlx/nn/layers/quantized.pyi
+++ b/.mlx_typings/mlx/nn/layers/quantized.pyi
@@ -2,7 +2,7 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import mlx.core as mx
 from base import Module
@@ -13,8 +13,10 @@ def quantize(
     bits: int = ...,
     *,
     mode: str = ...,
-    class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
-):  # -> None:
+    class_predicate: Optional[
+        Callable[[str, Module], Union[bool, dict[str, Any]]]
+    ] = ...,
+) -> None:
     """Quantize the sub-modules of a module according to a predicate.
 
     By default all layers that define a ``to_quantized(group_size, bits)``
diff --git a/.mlx_typings/mlx_lm/models/gemma4.pyi b/.mlx_typings/mlx_lm/models/gemma4.pyi
new file mode 100644
index 00000000..0ff6ef87
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/gemma4.pyi
@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import gemma4_text
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: Optional[dict[str, Any]]
+    vocab_size: int
+
+    def __post_init__(self) -> None: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: gemma4_text.Model
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[gemma4_text.DecoderLayer]: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...
diff --git a/.mlx_typings/mlx_lm/models/gemma4_text.pyi b/.mlx_typings/mlx_lm/models/gemma4_text.pyi
new file mode 100644
index 00000000..728d91c1
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/gemma4_text.pyi
@@ -0,0 +1,179 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    head_dim: int
+    global_head_dim: int
+    global_partial_rotary_factor: float
+    rms_norm_eps: float
+    vocab_size: int
+    vocab_size_per_layer_input: int
+    num_key_value_heads: int
+    num_global_key_value_heads: Optional[int]
+    num_kv_shared_layers: int
+    pad_token_id: int
+    hidden_size_per_layer_input: int
+    rope_traditional: bool
+    partial_rotary_factor: float
+    rope_parameters: Optional[Dict[str, Any]]
+    sliding_window: int
+    sliding_window_pattern: int
+    max_position_embeddings: int
+    attention_k_eq_v: bool
+    final_logit_softcapping: float
+    use_double_wide_mlp: bool
+    enable_moe_block: bool
+    num_experts: Optional[int]
+    top_k_experts: Optional[int]
+    moe_intermediate_size: Optional[int]
+    layer_types: Optional[List[str]]
+    tie_word_embeddings: bool
+
+    def __post_init__(self) -> None: ...
+
+class MLP(nn.Module):
+    gate_proj: nn.Linear
+    down_proj: nn.Linear
+    up_proj: nn.Linear
+
+    def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Router(nn.Module):
+    proj: nn.Linear
+    scale: mx.array
+    per_expert_scale: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class Experts(nn.Module):
+    switch_glu: SwitchGLU
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+    ) -> mx.array: ...
+
+class Attention(nn.Module):
+    layer_idx: int
+    layer_type: str
+    is_sliding: bool
+    head_dim: int
+    n_heads: int
+    n_kv_heads: int
+    use_k_eq_v: bool
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    v_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+class DecoderLayer(nn.Module):
+    layer_idx: int
+    layer_type: str
+    self_attn: Attention
+    mlp: MLP
+    enable_moe: bool
+    router: Router
+    experts: Experts
+    input_layernorm: nn.Module
+    post_attention_layernorm: nn.Module
+    pre_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm_1: nn.Module
+    post_feedforward_layernorm_2: nn.Module
+    pre_feedforward_layernorm_2: nn.Module
+    hidden_size_per_layer_input: int
+    per_layer_input_gate: Optional[nn.Linear]
+    per_layer_projection: Optional[nn.Linear]
+    post_per_layer_input_norm: Optional[nn.Module]
+    layer_scalar: mx.array
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = ...,
+        cache: Optional[Any] = ...,
+        per_layer_input: Optional[mx.array] = ...,
+        shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
+        offset: Optional[mx.array] = ...,
+    ) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...
+
+class Gemma4TextModel(nn.Module):
+    config: ModelArgs
+    vocab_size: int
+    window_size: int
+    sliding_window_pattern: int
+    num_hidden_layers: int
+    embed_tokens: nn.Embedding
+    embed_scale: float
+    layers: list[DecoderLayer]
+    norm: nn.Module
+    hidden_size_per_layer_input: int
+    embed_tokens_per_layer: Optional[nn.Embedding]
+    per_layer_model_projection: Optional[nn.Linear]
+    per_layer_projection_norm: Optional[nn.Module]
+    previous_kvs: list[int]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: Optional[mx.array] = ...,
+        cache: Optional[list[Any]] = ...,
+        input_embeddings: Optional[mx.array] = ...,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _get_per_layer_inputs(
+        self,
+        input_ids: Optional[mx.array],
+        input_embeddings: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _project_per_layer_inputs(
+        self,
+        input_embeddings: mx.array,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: Gemma4TextModel
+    final_logit_softcapping: float
+    tie_word_embeddings: bool
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    @property
+    def head_dim(self) -> int: ...
+    @property
+    def n_kv_heads(self) -> int: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...
diff --git a/.mlx_typings/mlx_lm/tokenizer_utils.pyi b/.mlx_typings/mlx_lm/tokenizer_utils.pyi
index 1326f059..b0ed12c0 100644
--- a/.mlx_typings/mlx_lm/tokenizer_utils.pyi
+++ b/.mlx_typings/mlx_lm/tokenizer_utils.pyi
@@ -117,6 +117,8 @@ class TokenizerWrapper:
     think_end: str | None
     think_start_id: int | None
     think_end_id: int | None
+    think_start_tokens: list[int] | None
+    think_end_tokens: list[int] | None
 
     def __init__(
         self,
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml
new file mode 100644
index 00000000..a36497e8
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-26b-a4b-it-4bit"
+n_layers = 30
+hidden_size = 2816
+num_key_value_heads = 8
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "4bit"
+base_model = "Gemma 4 26B A4B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 15608614044
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml
new file mode 100644
index 00000000..bd582311
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-26b-a4b-it-6bit"
+n_layers = 30
+hidden_size = 2816
+num_key_value_heads = 8
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "6bit"
+base_model = "Gemma 4 26B A4B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 21781015708
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml
new file mode 100644
index 00000000..9dda6aa4
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-26b-a4b-it-8bit"
+n_layers = 30
+hidden_size = 2816
+num_key_value_heads = 8
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "8bit"
+base_model = "Gemma 4 26B A4B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 27953417372
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-bf16.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-bf16.toml
new file mode 100644
index 00000000..ece39982
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-bf16.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-26b-a4b-it-bf16"
+n_layers = 30
+hidden_size = 2816
+num_key_value_heads = 8
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "bf16"
+base_model = "Gemma 4 26B A4B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 51611872412
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-31b-it-4bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-4bit.toml
new file mode 100644
index 00000000..d8a40f35
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-4bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-31b-it-4bit"
+n_layers = 60
+hidden_size = 5376
+num_key_value_heads = 16
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "4bit"
+base_model = "Gemma 4 31B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 18411755224
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-31b-it-6bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-6bit.toml
new file mode 100644
index 00000000..6222ce70
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-6bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-31b-it-6bit"
+n_layers = 60
+hidden_size = 5376
+num_key_value_heads = 16
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "6bit"
+base_model = "Gemma 4 31B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 26087306968
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-31b-it-8bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-8bit.toml
new file mode 100644
index 00000000..86396118
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-8bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-31b-it-8bit"
+n_layers = 60
+hidden_size = 5376
+num_key_value_heads = 16
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "8bit"
+base_model = "Gemma 4 31B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 33762858712
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-31b-it-bf16.toml b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-bf16.toml
new file mode 100644
index 00000000..1d4f740c
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-31b-it-bf16.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-31b-it-bf16"
+n_layers = 60
+hidden_size = 5376
+num_key_value_heads = 16
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "bf16"
+base_model = "Gemma 4 31B"
+capabilities = ["text", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 62546177752
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-4bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-4bit.toml
new file mode 100644
index 00000000..9d8f9932
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-4bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e2b-it-4bit"
+n_layers = 35
+hidden_size = 1536
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "4bit"
+base_model = "Gemma 4 E2B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 3580765126
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-6bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-6bit.toml
new file mode 100644
index 00000000..c9c6d1a0
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-6bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e2b-it-6bit"
+n_layers = 35
+hidden_size = 1536
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "6bit"
+base_model = "Gemma 4 E2B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 4739998662
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-8bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-8bit.toml
new file mode 100644
index 00000000..804da907
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-8bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e2b-it-8bit"
+n_layers = 35
+hidden_size = 1536
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "8bit"
+base_model = "Gemma 4 E2B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 5899232198
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-bf16.toml b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-bf16.toml
new file mode 100644
index 00000000..8f2920da
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e2b-it-bf16.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e2b-it-bf16"
+n_layers = 35
+hidden_size = 1536
+num_key_value_heads = 1
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "bf16"
+base_model = "Gemma 4 E2B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 10246357958
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-4bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-4bit.toml
new file mode 100644
index 00000000..1122bbeb
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-4bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e4b-it-4bit"
+n_layers = 42
+hidden_size = 2560
+num_key_value_heads = 2
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "4bit"
+base_model = "Gemma 4 E4B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 5216992212
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-6bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-6bit.toml
new file mode 100644
index 00000000..6f3b430c
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-6bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e4b-it-6bit"
+n_layers = 42
+hidden_size = 2560
+num_key_value_heads = 2
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "6bit"
+base_model = "Gemma 4 E4B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 7090961364
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-8bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-8bit.toml
new file mode 100644
index 00000000..48e21fef
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-8bit.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e4b-it-8bit"
+n_layers = 42
+hidden_size = 2560
+num_key_value_heads = 2
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "8bit"
+base_model = "Gemma 4 E4B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 8964930516
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-bf16.toml b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-bf16.toml
new file mode 100644
index 00000000..db87c2ee
--- /dev/null
+++ b/resources/inference_model_cards/mlx-community--gemma-4-e4b-it-bf16.toml
@@ -0,0 +1,15 @@
+model_id = "mlx-community/gemma-4-e4b-it-bf16"
+n_layers = 42
+hidden_size = 2560
+num_key_value_heads = 2
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "gemma"
+quantization = "bf16"
+base_model = "Gemma 4 E4B"
+capabilities = ["text", "vision"]
+
+context_length = 131072
+
+[storage_size]
+in_bytes = 15992314836
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 50da8125..bdbe3f63 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -261,6 +261,7 @@ class ConfigData(BaseModel):
             ["GptOssForCausalLM"],
             ["Step3p5ForCausalLM"],
             ["NemotronHForCausalLM"],
+            ["Gemma4ForConditionalGeneration"],
         ]
 
     @model_validator(mode="before")
@@ -287,7 +288,7 @@ class ConfigData(BaseModel):
         image_token_id = data.get("image_token_id")
         if vision_config is not None and image_token_id is not None:
             model_type = str(
-                vision_config.get("model_type", data.get("model_type", ""))  # pyright: ignore[reportAny]
+                data.get("model_type", vision_config.get("model_type", ""))  # pyright: ignore[reportAny]
             )
             assert info.context is not None
 
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 0085e2f5..3b6e8421 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -447,7 +447,7 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         )
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
-        if cache is not None:
+        if cache is not None and len(cache) > 0:  # type: ignore
             last = cache[-1]  # type: ignore
             dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
             if hasattr(dep_cache, "keys") and dep_cache.keys is not None:  # type: ignore
diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index 6376917d..ad941f90 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -338,7 +338,7 @@ def _entry_length(
 
 def cache_length(cache: KVCacheType) -> int:
     """Get the number of tokens in a KV cache."""
-    return max(_entry_length(c) for c in cache)
+    return max((_entry_length(c) for c in cache), default=0)
 
 
 def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 1d47ad40..61b927a9 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -187,7 +187,11 @@ class ExoBatchGenerator:
 
         vision_ctx = (
             patch_embed_tokens(
-                self.model, vision.embeddings, prefix_hit_length, len(prompt_tokens) - 1
+                self.model,
+                vision.embeddings,
+                prefix_hit_length,
+                len(prompt_tokens) - 1,
+                image_token_id=vision.image_token_id,
             )
             if vision is not None
             else contextlib.nullcontext()
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 3c0cb2ac..3a439f3e 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -77,7 +77,11 @@ _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
 
 @contextlib.contextmanager
 def patch_embed_tokens(
-    model: Model, embeddings: mx.array, start_offset: int = 0, token_count: int = 0
+    model: Model,
+    embeddings: mx.array,
+    start_offset: int = 0,
+    token_count: int = 0,
+    image_token_id: int | None = None,
 ) -> Generator[None]:
     inner = get_inner_model(model)  # type: ignore
     original_embed = inner.embed_tokens  # type: ignore
@@ -85,15 +89,30 @@ def patch_embed_tokens(
     offset = [start_offset]
 
     def _inject(input_ids: mx.array) -> mx.array:
-        start = offset[0]
-        if start >= end_offset:
-            return original_embed(input_ids)  # type: ignore
+        chunk_start = offset[0]
         chunk_len = input_ids.shape[-1]
-        end = min(start + chunk_len, end_offset)
-        offset[0] = end
-        if end - start < chunk_len:
+        chunk_end = chunk_start + chunk_len
+        offset[0] = chunk_end
+
+        # The injection window is [start_offset, end_offset).
+        if chunk_end <= start_offset or chunk_start >= end_offset:
             return original_embed(input_ids)  # type: ignore
-        return embeddings[:, start:end, :]
+
+        # Mixed chunk: splice the pre-computed embeddings for the overlap
+        # into `original_embed(input_ids)` for any text-only fringes.
+        overlap_start = max(chunk_start, start_offset)
+        overlap_end = min(chunk_end, end_offset)
+        dst_start = overlap_start - chunk_start
+        dst_end = overlap_end - chunk_start
+        text_embeds: mx.array = original_embed(input_ids)  # type: ignore
+        return mx.concatenate(
+            [
+                text_embeds[:, :dst_start, :],
+                embeddings[:, overlap_start:overlap_end, :],
+                text_embeds[:, dst_end:, :],
+            ],
+            axis=1,
+        )
 
     for attr in dir(original_embed):  # type: ignore
         if not attr.startswith("_") and not hasattr(_inject, attr):
@@ -101,10 +120,30 @@ def patch_embed_tokens(
                 setattr(_inject, attr, getattr(original_embed, attr))  # type: ignore
 
     inner.embed_tokens = _inject
+
+    # Gemma 4 (e2b/e4b) has a second, independent embedding table that produces
+    # per-layer conditioning signals via self.embed_tokens_per_layer(input_ids).
+    # The injected vision embeddings live in the main residual stream only, so
+    # if image_token_id positions are passed through as-is the per-layer table
+    # produces garbage signals at those positions (the `<image>` token was never
+    # trained to have meaningful per-layer inputs).
+    original_per_layer = getattr(inner, "embed_tokens_per_layer", None)  # type: ignore
+    if original_per_layer is not None and image_token_id is not None:
+
+        def _clean_per_layer(input_ids: mx.array) -> mx.array:
+            clean_ids = mx.where(
+                input_ids == image_token_id, mx.zeros_like(input_ids), input_ids
+            )
+            return original_per_layer(clean_ids)  # type: ignore
+
+        inner.embed_tokens_per_layer = _clean_per_layer
+
     try:
         yield
     finally:
         inner.embed_tokens = original_embed
+        if original_per_layer is not None and image_token_id is not None:
+            inner.embed_tokens_per_layer = original_per_layer
 
 
 class PrefillCancelled(BaseException):
@@ -576,7 +615,11 @@ def mlx_generate(
 
     maybe_vision_ctx = (
         patch_embed_tokens(
-            model, vision.embeddings, prefix_hit_length, len(prompt_tokens) - 1
+            model,
+            vision.embeddings,
+            prefix_hit_length,
+            len(prompt_tokens) - 1,
+            image_token_id=vision.image_token_id,
         )
         if vision is not None
         else contextlib.nullcontext()
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index b4eca6aa..63f1b6ed 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -336,6 +336,8 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
     elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower:
         # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>)
         return [248046, 248044]
+    elif "gemma-4" in model_id_lower or "gemma-3" in model_id_lower:
+        return [1, 106, 50]
     return None
 
 
@@ -407,23 +409,13 @@ def load_tokenizer_for_model_id(
             tool_parser=_parse_kimi_tool_calls,
         )
 
+    # We should really consider going back to mlx lm load to get tokenizer
     tokenizer = load_tokenizer(
         model_path,
         tokenizer_config_extra={"trust_remote_code": trust_remote_code},
         eos_token_ids=eos_token_ids,
     )
 
-    if "gemma-3" in model_id_lower:
-        gemma_3_eos_id = 1
-        gemma_3_end_of_turn_id = 106
-        if tokenizer.eos_token_ids is not None:
-            if gemma_3_end_of_turn_id not in tokenizer.eos_token_ids:
-                tokenizer.eos_token_ids = list(tokenizer.eos_token_ids) + [
-                    gemma_3_end_of_turn_id
-                ]
-        else:
-            tokenizer.eos_token_ids = [gemma_3_eos_id, gemma_3_end_of_turn_id]
-
     return tokenizer
 
 
@@ -671,21 +663,37 @@ def fix_unmatched_think_end_tokens(
 ) -> mx.array:
     if not tokenizer.has_thinking:
         return tokens
-    assert tokenizer.think_start_id
-    assert tokenizer.think_end_id
-    think_start_id: int = tokenizer.think_start_id
-    think_end_id: int = tokenizer.think_end_id
+    assert tokenizer.think_start_tokens
+    assert tokenizer.think_end_tokens
+    think_start_tokens: list[int] = tokenizer.think_start_tokens
+    think_end_tokens: list[int] = tokenizer.think_end_tokens
     token_list: list[int] = cast(list[int], tokens.tolist())
     result: list[int] = []
+
     depth = 0
+    accumulated_think_start_length = 0
+    accumulated_think_end_length = 0
+
     for token in token_list:
-        if token == think_start_id:
-            depth += 1
-        elif token == think_end_id:
-            if depth == 0:
-                result.append(think_start_id)
-            else:
-                depth -= 1
+        if token == think_start_tokens[accumulated_think_start_length]:
+            accumulated_think_start_length += 1
+            if accumulated_think_start_length == len(think_start_tokens):
+                depth += 1
+                accumulated_think_start_length = 0
+
+        elif token == think_end_tokens[accumulated_think_end_length]:
+            accumulated_think_end_length += 1
+            if accumulated_think_end_length == len(think_end_tokens):
+                if depth == 0:
+                    result.extend(think_start_tokens)
+                else:
+                    depth -= 1
+                accumulated_think_end_length = 0
+
+        else:
+            accumulated_think_start_length = 0
+            accumulated_think_end_length = 0
+
         result.append(token)
     return mx.array(result)
 
diff --git a/src/exo/worker/engines/mlx/vision.py b/src/exo/worker/engines/mlx/vision.py
index b2cc8ea4..10d5312f 100644
--- a/src/exo/worker/engines/mlx/vision.py
+++ b/src/exo/worker/engines/mlx/vision.py
@@ -6,8 +6,9 @@ import inspect
 import io
 import json
 import re
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 if TYPE_CHECKING:
     from mlx_vlm.utils import ImageProcessor
@@ -33,13 +34,64 @@ from exo.worker.engines.mlx.utils_mlx import (
 )
 from exo.worker.runner.bootstrap import logger
 
+_video_processor_patched = False
+
 
 def _filter_config(cls: type, d: dict[str, Any]) -> dict[str, Any]:
     valid = set(inspect.signature(cls.__init__).parameters.keys()) - {"self"}
     return {k: v for k, v in d.items() if k in valid}  # type: ignore
 
 
-_video_processor_patched = False
+_ProcessorOutput = dict[str, np.ndarray] | tuple[dict[str, np.ndarray], list[int]]
+
+
+def _run_processor(
+    processor: "ImageProcessor",
+    pil_images: list[Image.Image],
+) -> tuple[dict[str, np.ndarray], list[int] | None]:
+    """
+    Image processors split into two families by how they report per-image
+    token counts:
+
+    1. Variable-resolution patch models (Qwen3-VL, Llama 4 vision, ...):
+       return a `BatchFeature` dict containing `pixel_values` and
+       `image_grid_thw` — an (n_images, 3) array of (temporal, height, width)
+       patch counts. The caller multiplies the three to get the per-image
+       token count, so no override is needed.
+
+    2. Fixed-token-budget models (Gemma 4): every image collapses to a fixed
+       number of soft tokens, so there's no grid to report. These processors
+       return `(batch_feature_dict, [n_tokens_per_image])` instead.
+
+    We normalize both into (dict, optional tokens override).
+    """
+    raw = cast(_ProcessorOutput, processor(images=pil_images, return_tensors="np"))
+    if isinstance(raw, tuple):
+        batch, tokens = raw
+        return batch, [int(n) for n in tokens]
+    return raw, None
+
+
+def _instantiate_projector(
+    cls: type,
+    model_config: Any,  # pyright: ignore[reportAny]
+    vision_config: Any,  # pyright: ignore[reportAny]
+    text_config: Any,  # pyright: ignore[reportAny]
+) -> nn.Module:
+    """
+    Instantiate projector/embedding classes with any missing values
+    """
+    init_sig = inspect.signature(cls.__init__)
+    params = {n: p for n, p in init_sig.parameters.items() if n != "self"}
+    kwargs: dict[str, Any] = {}
+
+    if "embedding_dim" in params:
+        kwargs["embedding_dim"] = vision_config.hidden_size  # pyright: ignore[reportAny]
+    if "text_hidden_size" in params:
+        kwargs["text_hidden_size"] = text_config.hidden_size  # pyright: ignore[reportAny]
+    if "eps" in params:
+        kwargs["eps"] = getattr(vision_config, "rms_norm_eps", 1e-6)  # pyright: ignore[reportAny]
+    return cls(**kwargs)  # type: ignore
 
 
 def _patch_video_processor() -> None:
@@ -128,6 +180,7 @@ class VisionResult:
     prompt_tokens: mx.array
     embeddings: mx.array
     media_regions: list[MediaRegion]
+    image_token_id: int
 
 
 class VisionEncoder:
@@ -159,6 +212,45 @@ class VisionEncoder:
             results.append(importlib.import_module(name))
         return results[0] if len(results) == 1 else tuple(results)
 
+    def _apply_projector_quantization_if_needed(
+        self, projector_weights: dict[str, mx.array]
+    ) -> None:
+        # Quantized models ship the projector's Linear layers as packed uint32
+        # weights plus `.scales`/`.biases`. Our now instantiated projector
+        # uses plain nn.Linear, so we must mirror the packing via nn.quantize
+        # before load_weights, otherwise MLX rejects the extra parameters.
+        if self._projector is None:
+            return
+        has_quantized_tensors = any(
+            key.endswith((".scales", ".biases")) or val.dtype == mx.uint32
+            for key, val in projector_weights.items()
+        )
+        if not has_quantized_tensors:
+            return
+        config = self._load_config_json()
+        quant_cfg = cast(dict[str, Any], config.get("quantization") or {})
+        if not quant_cfg:
+            return
+        group_size = int(cast(int, quant_cfg.get("group_size", 64)))
+        bits = int(cast(int, quant_cfg.get("bits", 4)))
+        nn.quantize(self._projector, group_size=group_size, bits=bits)
+
+    def _load_image_processor_from_module(self, repo: str) -> "ImageProcessor | None":
+        # mlx_vlm.utils.load_image_processor only works for models that set
+        # `Model.ImageProcessor = <cls>`, but Gemma4 just uses
+        # `Gemma4ImageProcessor` from the package `__init__.py`
+        try:
+            pkg: Any = importlib.import_module(
+                f"mlx_vlm.models.{self._config.model_type}"
+            )
+        except ImportError:
+            return None
+        for attr in dir(pkg):  # pyright: ignore[reportAny]
+            cls = getattr(pkg, attr)  # pyright: ignore[reportAny]
+            if isinstance(cls, type) and attr.endswith("ImageProcessor"):
+                return cls.from_pretrained(repo)  # type: ignore
+        return None
+
     def ensure_loaded(self) -> None:
         if self._loaded:
             return
@@ -194,7 +286,7 @@ class VisionEncoder:
                 if (
                     isinstance(obj, type)
                     and issubclass(obj, nn.Module)
-                    and "Projector" in attr_name
+                    and ("Projector" in attr_name or "Embedder" in attr_name)
                 ):
                     projector_cls = obj
                     break
@@ -214,7 +306,12 @@ class VisionEncoder:
                 vision_config=vision_config,
                 **_filter_config(config_mod.ModelConfig, extra),  # type: ignore
             )
-            self._projector = projector_cls(model_config)  # type: ignore
+            self._projector = _instantiate_projector(
+                projector_cls,
+                model_config,
+                vision_config,
+                text_config,
+            )
 
         processor_repo = self._config.processor_repo
         if processor_repo:
@@ -223,7 +320,9 @@ class VisionEncoder:
             self._load_weights_from_model_repo()
 
         repo = processor_repo or str(self._model_path)
-        image_proc = load_image_processor(repo)
+        image_proc = load_image_processor(
+            repo
+        ) or self._load_image_processor_from_module(repo)
         if image_proc is not None:
             self._processor = image_proc
         else:
@@ -295,18 +394,32 @@ class VisionEncoder:
             raise FileNotFoundError(f"No safetensors files found in {self._model_path}")
 
         vision_prefixes = ["vision_tower.", "model.visual."]
+        projector_prefixes = [
+            "embed_vision.",
+            "multi_modal_projector.",
+            "mm_projector.",
+        ]
         vision_weights: dict[str, mx.array] = {}
-        found_raw_prefix = False
+        projector_weights: dict[str, mx.array] = {}
+
+        # If weights under `model.visual.`, we need to call mlx_vlm's VisionModel.sanitize()
+        # to remap into its own keys.
+        needs_sanitize = False
+
         for sf_path in safetensors_files:
             file_weights: dict[str, mx.array] = mx.load(str(sf_path))  # type: ignore
             for key, val in file_weights.items():
                 for prefix in vision_prefixes:
                     if key.startswith(prefix):
-                        short_key = key[len(prefix) :]
-                        vision_weights[short_key] = val
+                        vision_weights[key[len(prefix) :]] = val
                         if prefix == "model.visual.":
-                            found_raw_prefix = True
+                            needs_sanitize = True
                         break
+                else:
+                    for prefix in projector_prefixes:
+                        if key.startswith(prefix):
+                            projector_weights[key[len(prefix) :]] = val
+                            break
 
         if not vision_weights:
             raise ValueError(
@@ -315,14 +428,27 @@ class VisionEncoder:
             )
 
         assert self._vision_tower is not None
-        if found_raw_prefix and hasattr(self._vision_tower, "sanitize"):
-            vision_weights = self._vision_tower.sanitize(vision_weights)  # type: ignore
+        if needs_sanitize:
+            sanitize: Callable[[dict[str, mx.array]], dict[str, mx.array]] | None = (
+                getattr(self._vision_tower, "sanitize", None)
+            )
+            if sanitize is not None:
+                vision_weights = sanitize(vision_weights)
 
-        self._vision_tower.load_weights(list(vision_weights.items()))  # type: ignore
+        self._vision_tower.load_weights(list(vision_weights.items()))
         mx.eval(self._vision_tower.parameters())
 
-        n_vision = sum(v.size for _, v in vision_weights.items())  # type: ignore
-        logger.info(f"Vision encoder loaded: {n_vision / 1e6:.1f}M params")
+        if self._projector is not None and projector_weights:
+            self._apply_projector_quantization_if_needed(projector_weights)
+            self._projector.load_weights(list(projector_weights.items()))
+            mx.eval(self._projector.parameters())
+
+        n_vision = sum(v.size for v in vision_weights.values())
+        n_proj = sum(v.size for v in projector_weights.values())
+        logger.info(
+            f"Vision encoder loaded: {n_vision / 1e6:.1f}M params"
+            + (f" (+ projector {n_proj / 1e6:.1f}M)" if n_proj else "")
+        )
 
     def encode_images(self, images: list[Base64Image]) -> tuple[mx.array, list[int]]:
         self.ensure_loaded()
@@ -333,12 +459,21 @@ class VisionEncoder:
         for idx, img in enumerate(pil_images):
             logger.info(f"Image {idx}: {img.width}x{img.height} mode={img.mode}")
 
+        per_image_pixels: list[mx.array]
+        grid_thw: mx.array | None
+        n_tokens_per_image: list[int]
+
         if self._config.processor_repo:
             processed = self._processor.preprocess(
                 [{"type": "image", "image": img} for img in pil_images],
                 return_tensors="np",
             )
-            pixel_values = mx.array(processed["pixel_values"])  # type: ignore
+            stacked_pixels = mx.array(processed["pixel_values"])  # type: ignore
+            if stacked_pixels.ndim == 3:
+                stacked_pixels = stacked_pixels[None]
+            per_image_pixels = [
+                stacked_pixels[i : i + 1] for i in range(stacked_pixels.shape[0])
+            ]
             grid_thw = mx.array(processed["grid_thws"])  # type: ignore
             assert self._merge_kernel_size is not None
             merge_length = int(np.prod(self._merge_kernel_size))
@@ -347,31 +482,75 @@ class VisionEncoder:
                 for i in range(grid_thw.shape[0])
             ]
         else:
-            processed = self._processor(
-                images=pil_images,
-                return_tensors="np",
+            batch, tokens_override = _run_processor(self._processor, pil_images)
+            # `Gemma4ImageProcessor` returns pixel_values as a plain ndarray
+            # when all images resize to the same shape, or as a Python list of
+            # per-image (C, H, W) ndarrays when they differ. Treat it as the
+            # union here.
+            raw_pixel_values: np.ndarray | list[np.ndarray] = cast(
+                "np.ndarray | list[np.ndarray]", batch["pixel_values"]
             )
-            pixel_values = mx.array(processed["pixel_values"])  # type: ignore
-            grid_thw = mx.array(processed["image_grid_thw"])  # type: ignore
-            merge_unit = self._spatial_merge_size**2
-            n_tokens_per_image = [
-                int(
-                    grid_thw[i, 0].item()
-                    * grid_thw[i, 1].item()
-                    * grid_thw[i, 2].item()
-                )
-                // merge_unit
-                for i in range(grid_thw.shape[0])
-            ]
+            raw_grid = batch.get("image_grid_thw")
+            grid_thw = mx.array(raw_grid) if raw_grid is not None else None
+            if tokens_override is not None:
+                n_tokens_per_image = tokens_override
+            else:
+                assert grid_thw is not None
+                merge_unit = self._spatial_merge_size**2
+                n_tokens_per_image = [
+                    int(
+                        grid_thw[i, 0].item()
+                        * grid_thw[i, 1].item()
+                        * grid_thw[i, 2].item()
+                    )
+                    // merge_unit
+                    for i in range(grid_thw.shape[0])
+                ]
+
+            if isinstance(raw_pixel_values, list):
+                per_image_pixels = [
+                    # (C, H, W) -> (1, C, H, W)
+                    mx.array(p)[None] if p.ndim == 3 else mx.array(p)
+                    for p in raw_pixel_values
+                ]
+            else:
+                stacked = mx.array(raw_pixel_values)
+                per_image_pixels = [stacked[i : i + 1] for i in range(stacked.shape[0])]
+
+        patch_embed_weight = None
+        with contextlib.suppress(AttributeError):
+            patch_embed_weight = self._vision_tower.patch_embed.proj.weight  # type: ignore
+        with contextlib.suppress(AttributeError):
+            patch_embed_weight = self._vision_tower.patch_embedder.input_proj.weight  # type: ignore
+        assert patch_embed_weight is not None, (
+            "vision tower has no recognised patch-embedding linear"
+        )
+        tower_dtype = cast(mx.Dtype, patch_embed_weight.dtype)
 
         if self._needs_nhwc:
+            assert grid_thw is not None
+            pixel_values = mx.concatenate(per_image_pixels, axis=0).astype(tower_dtype)
             grid_hw = grid_thw[:, 1:] if grid_thw.shape[-1] == 3 else grid_thw
             hidden_states = self._vision_tower(
                 pixel_values.transpose(0, 2, 3, 1),
                 output_hidden_states=True,
                 grid_thw=grid_hw,
             )
+        elif grid_thw is None:
+            # Fixed-token-budget models (gemma4): run each image separately
+            # since they can have different spatial shapes *and* different
+            # soft-token counts, then flatten each to (n_tokens_i, hidden)
+            # and concatenate along the token axis.
+            per_image_hidden: list[mx.array] = []
+            for pv in per_image_pixels:
+                result = self._vision_tower(pv.astype(tower_dtype))
+                h = result[0] if isinstance(result, tuple) else result
+                if h.ndim == 3:
+                    h = h.reshape(-1, h.shape[-1])
+                per_image_hidden.append(h)
+            hidden_states = mx.concatenate(per_image_hidden, axis=0)
         else:
+            pixel_values = mx.concatenate(per_image_pixels, axis=0).astype(tower_dtype)
             result = self._vision_tower(pixel_values, grid_thw)
             hidden_states = result[0] if isinstance(result, tuple) else result
 
@@ -380,6 +559,11 @@ class VisionEncoder:
         else:
             image_features = hidden_states
 
+        # `create_vision_embeddings` expects a 2D (total_tokens, hidden) view,
+        # but fixed-token-budget models (gemma4) return (n_images, tokens, hidden).
+        if image_features.ndim == 3:
+            image_features = image_features.reshape(-1, image_features.shape[-1])
+
         return image_features, n_tokens_per_image
 
 
@@ -420,6 +604,16 @@ def create_vision_embeddings(
             n = min(n_placeholders, image_features.shape[0])
             image_features = image_features[:n]
 
+        # Gemma-family models apply `h = input_embeddings * embed_scale` inside
+        # the inner model's forward pass. That scale is appropriate for text
+        # tokens (which come out of a raw `embed_tokens(id)` lookup) but not
+        # for our pre-projected image features. Pre-divide by `embed_scale`
+        # so that after the model multiplies, image features are unchanged
+        # while text positions remain correctly scaled.
+        if hasattr(inner, "embed_scale"):  # type: ignore
+            embed_scale = float(inner.embed_scale)  # type: ignore
+            image_features = image_features / embed_scale
+
         image_indices = mx.cumsum(is_image.astype(mx.int32)) - 1
         image_indices = mx.clip(image_indices, 0, image_features.shape[0] - 1)
 
@@ -563,6 +757,7 @@ class VisionProcessor:
             prompt_tokens=prompt_tokens,
             embeddings=embeddings,
             media_regions=media_regions,
+            image_token_id=self.vision_config.image_token_id,
         )
 
 
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 8cdc3ffc..573afb76 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -339,7 +339,9 @@ class Worker:
                             f"from {len(chunk_buffer)} chunks"
                         )
 
-                    resolved_images = [by_index[i] for i in sorted(by_index)]
+                    resolved_images = [
+                        Base64Image(by_index[i]) for i in sorted(by_index)
+                    ]
                     modified_task = task.model_copy(
                         update={
                             "task_params": task.task_params.model_copy(
diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 9f380a77..c7e50649 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -329,21 +329,35 @@ def parse_thinking_models(
     Always yields tokens with finish_reason to avoid hanging the chunk stream.
     """
     is_thinking = starts_in_thinking
+    accumulated = ""
+
     for response in responses:
         if response is None:
             yield None
             continue
+
+        accumulated += response.text
+
         if response.finish_reason is not None:
             yield response.model_copy(update={"is_thinking": False})
             continue
 
-        if response.text == think_start:
+        if accumulated == think_start and not is_thinking:
             is_thinking = True
+            accumulated = ""
             continue
-        if response.text == think_end:
+        if accumulated == think_end and is_thinking:
             is_thinking = False
+            accumulated = ""
+            continue
+
+        if (think_start and accumulated == think_start[: len(accumulated)]) or (
+            think_end and accumulated == think_end[: len(accumulated)]
+        ):
             continue
 
+        accumulated = ""
+
         yield response.model_copy(update={"is_thinking": is_thinking})
 
 
diff --git a/uv.lock b/uv.lock
index f40615ff..850dbe0b 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1541,7 +1541,7 @@ dependencies = [
 
 [[package]]
 name = "mlx-vlm"
-version = "0.4.1"
+version = "0.4.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "datasets", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1558,9 +1558,9 @@ dependencies = [
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "uvicorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/57/8f/31204f1a8c7404e523a5578949ea668e668e10dd67a0f63336f261014c0f/mlx_vlm-0.4.1.tar.gz", hash = "sha256:4e2d8a232715dbca72d346f43cf54d5738452848855792ffb1b285228ae7c7bd", size = 621840, upload-time = "2026-03-21T14:26:04.586Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/ec/108aec30efb159940ea29d133d5d8ec14840edbec914869b46eaafac5552/mlx_vlm-0.4.4.tar.gz", hash = "sha256:3197e277c1be9ed1712ea04624df029e486f7747ad93e40e7bd1c9c771f8b179", size = 836370, upload-time = "2026-04-04T15:19:01.087Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c1/78/856f44c6bdd8791427fa59a093a1d00c91cdbe16238506602fc3968017bb/mlx_vlm-0.4.1-py3-none-any.whl", hash = "sha256:89feca2e8be31609770c0e8a6d88fa21d00ee25bd3d56b4aafce59d35dd63b71", size = 768806, upload-time = "2026-03-21T14:26:03.129Z" },
+    { url = "https://files.pythonhosted.org/packages/d9/81/235518176c3c8230e5274e91346ecf940591f653e73b0daeb505fb37eea9/mlx_vlm-0.4.4-py3-none-any.whl", hash = "sha256:3ff86ea738ab1914dc1b07e4fa5d4cc34bec5909e540692cfad0af808af13c11", size = 1014936, upload-time = "2026-04-04T15:18:59.328Z" },
 ]
 
 [[package]]

← 6172617b add env override to macos app (#1869)  ·  back to Exo  ·  build: remove pyinstaller temp artifacts (#1868) 4b13735e →