Core Layers of the MLX-VLM Architecture: A Deep Dive into the Multimodal Pipeline

MLX-VLM implements a three-stage multimodal pipeline consisting of a VisionModel tower for feature extraction, a LanguageModel backbone for generation, and a MultimodalEmbedder projection layer that bridges visual and text embeddings through special token insertion.

The Blaizzy/mlx-vlm repository provides a modular framework for running vision-language models on Apple Silicon using the MLX framework. Understanding the core layers of the MLX-VLM architecture is essential for customizing inference, debugging multimodal behavior, or extending the library to new model families. Each component is designed to be swappable, allowing different vision encoders and language models to plug into a standardized projection and attention infrastructure.

The Three-Stage Multimodal Pipeline

The architecture cleanly separates visual processing, language processing, and the fusion mechanism that connects them. These are the foundational layers present across all model implementations in the repository.

VisionModel (Vision Tower)

The VisionModel handles feature extraction from raw pixel inputs using vision encoders such as SigLIP, ViT, or DINOv2. This layer receives pixel_values tensors and outputs visual embeddings that the language model can interpret.

In mlx_vlm/models/<model>/vision.py (exemplified by the Qwen3-VL implementation at mlx_vlm/models/qwen3_vl/vision.py), the vision tower processes images through multiple transformer blocks to produce hidden_states. These features may also include deep-stack visual embeddings for high-resolution image understanding.

LanguageModel (LLM Backbone)

The LanguageModel serves as the autoregressive backbone that generates text logits from token embeddings. It supports optional KV-cache quantization for memory-efficient inference on resource-constrained devices.

According to the source code in mlx_vlm/models/<model>/language.py, this layer handles the forward pass through the transformer stack, applying causal masking and rotary position embeddings. The language model remains agnostic to whether its inputs come from text tokens or projected visual features.

Multimodal Embedder and Projection Layer

The MultimodalEmbedder acts as the critical bridge between vision and language modalities. It projects vision embeddings into the hidden-size space of the language model and inserts them at special token positions (such as <image> or <audio> tokens).

As implemented in mlx_vlm/models/gemma4/gemma4.py, the MultimodalEmbedder class provides the embed_vision method for linear projection and utilizes the masked_scatter routine to merge visual features into the token stream. This layer ensures the LLM treats visual tokens as native elements of its input sequence.

Supporting Infrastructure

Beyond the three primary stages, several utility layers manage attention mechanics and input formatting.

Attention Mechanisms and KV-Cache

The Attention layer provides scaled-dot-product attention used by both vision and language components. For efficient long-context generation, the architecture supports TurboQuant KV-cache quantization.

The function scaled_dot_product_attention in mlx_vlm/models/base.py routes attention computations and optionally integrates with TurboQuantKVCache (defined in mlx_vlm/turboquant.py) to reduce memory footprint during autoregressive decoding. This quantization happens transparently when use_turboquant is enabled on the model.

Input Embedding Wrapper

The InputEmbeddingsFeatures dataclass in mlx_vlm/models/base.py packages combined embeddings, visual masks, and optional deep-stack features into a single object passed to the language model. This abstraction standardizes how multimodal inputs flow through the architecture regardless of specific model implementations.

How the Layers Interact

The core layers of the MLX-VLM architecture operate in a specific sequence during inference:

1. Vision Forward: VisionModel receives pixel_values and returns feature tensors plus any deep-stack visual embeddings.
2. Projection: MultimodalEmbedder linearly maps vision outputs to the language hidden-size via embed_vision.
3. Token Merging: The masked_scatter helper (or model-specific variants like merge_input_ids_with_image_features) inserts projected visual embeddings at special image-token positions within the token embedding matrix.
4. Language Forward: LanguageModel processes the combined token and visual embeddings through its attention stack, utilizing scaled_dot_product_attention and optionally TurboQuantKVCache for quantized decoding.
5. Output Generation: The generation loop in mlx_vlm/generate.py decodes logits into text, completing the multimodal pipeline.

Working with Core Layers

These practical examples demonstrate how to interact directly with the architectural components.

Loading and Running Inference

To load a multimodal model and execute image-plus-text generation:

from mlx_vlm import load, generate
from mlx_vlm.prompt_utils import apply_chat_template

# Load a Gemma-4 model that supports both vision and audio

model, processor = load("google/gemma-4-e4b-it")

# Build a prompt that expects an image description

prompt = apply_chat_template(
    processor,
    model.config,
    "Describe the scene in the image.",
    num_images=1,               # tells the template to insert an <image> token

)

# Provide the image path (can also be a PIL image)

image = ["path/to/photo.jpg"]

result = generate(
    model=model,
    processor=processor,
    prompt=prompt,
    image=image,
    max_tokens=200,
    temperature=0.7,
)

print(result)

Key files referenced: load in mlx_vlm/__init__.py, apply_chat_template in mlx_vlm/prompt_utils.py, and generate in mlx_vlm/generate.py.

Inspecting Visual Embedding Injection

To directly examine how visual embeddings are injected into the language model:

from mlx_vlm.models.gemma4.gemma4 import Model as Gemma4Model
from mlx_vlm.models.gemma4.config import ModelConfig
import mlx.core as mx

cfg = ModelConfig.from_pretrained("google/gemma-4-e4b-it")
model = Gemma4Model(cfg)

# Dummy image tensor (batch=1, H=224, W=224, C=3)

pixel_values = mx.random.normal([1, 224, 224, 3])

# Dummy input ids containing an <image> token (id = cfg.image_token_id)

input_ids = mx.array([[cfg.bos_token_id, cfg.image_token_id, cfg.eos_token_id]])

embeds = model.get_input_embeddings(input_ids=input_ids,
                                    pixel_values=pixel_values)

print("Shape of combined embeddings:", embeds.inputs_embeds.shape)

This exercises the VisionModel (inside self.vision_tower), MultimodalEmbedder, and masked_scatter layers.

Enabling TurboQuant KV-Cache

For large models, enable quantized KV-caching to reduce memory usage:

from mlx_vlm.models.qwen3_5.qwen3_5 import Model as Qwen3_5Model
from mlx_vlm.models.qwen3_5.config import ModelConfig
import mlx.core as mx

cfg = ModelConfig.from_pretrained("mlx-community/Qwen3.5-2B-Instruct")
model = Qwen3_5Model(cfg)

# Turn on TurboQuant for the KV cache

model = model.apply(lambda m: setattr(m, "use_turboquant", True) or m)

# Simple generation loop (omitted for brevity)

Relevant source: TurboQuantKVCache implementation in mlx_vlm/turboquant.py and the attention routing in scaled_dot_product_attention (see mlx_vlm/models/base.py).

Summary

• MLX-VLM separates concerns into three distinct stages: vision encoding, multimodal projection, and language generation.
• The VisionModel in mlx_vlm/models/<model>/vision.py extracts features from images using architectures like SigLIP or DINOv2.
• The MultimodalEmbedder projects visual embeddings into the language model's hidden space and inserts them at special token positions via masked_scatter.
• The LanguageModel in mlx_vlm/models/<model>/language.py handles autoregressive generation with support for TurboQuantKVCache quantization.
• InputEmbeddingsFeatures and scaled_dot_product_attention in mlx_vlm/models/base.py provide standardized interfaces for multimodal data flow and attention computation.

Frequently Asked Questions

What is the role of the MultimodalEmbedder in MLX-VLM?

The MultimodalEmbedder serves as the bridge between vision and language modalities. It linearly projects vision encoder outputs into the language model's embedding dimension and manages the insertion of these visual tokens at specific positions (such as <image> tokens) within the input sequence. This layer is implemented in model-specific files like mlx_vlm/models/gemma4/gemma4.py and utilizes the masked_scatter routine to merge features without disrupting the text token stream.

How does MLX-VLM handle special image tokens?

MLX-VLM uses special token IDs (such as cfg.image_token_id) to mark positions in the input sequence where visual embeddings should be inserted. During the forward pass, the MultimodalEmbedder identifies these positions and uses masked_scatter to replace the placeholder token embeddings with actual projected vision features. The apply_chat_template function in mlx_vlm/prompt_utils.py automatically inserts these special tokens when processing chat-style prompts with images.

Where is the vision encoder implemented in MLX-VLM?

The vision encoder is implemented in model-specific vision.py files located at mlx_vlm/models/<model>/vision.py. For example, the Qwen3-VL vision tower resides in mlx_vlm/models/qwen3_vl/vision.py, while other models like Gemma-4 and Granite follow the same pattern. These files define the VisionModel class that processes pixel_values into feature tensors consumable by the projection layers.

Can MLX-VLM use quantized KV-caches for inference?

Yes, MLX-VLM supports TurboQuant KV-cache quantization for memory-efficient inference. By setting use_turboquant to True on the model instance, the scaled_dot_product_attention function in mlx_vlm/models/base.py routes cache operations through the TurboQuantKVCache class defined in mlx_vlm/turboquant.py. This reduces memory footprint during long-context generation without requiring changes to the model architecture or loading procedure.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}