Implementation:Predibase Lorax MLlama Vision Model

Knowledge Sources	Predibase_Lorax
Domains	Model_Architecture, Inference
Last Updated	2026-02-08 00:00 GMT

Overview

Implements the MLlama (Meta Llama) vision-language model architecture, including the vision encoder with transformer and global transformer stages, cross-attention text layers, and the full multimodal forward pass for the LoRax inference server.

Description

This module contains the complete MLlama vision-language model implementation spanning 1096 lines. It defines both the vision encoder pipeline and the cross-attention text model integration.

Vision Components:

• MllamaVisionMLP (extends nn.Module) - Feed-forward network within vision encoder layers with LoRA adapter support via TensorParallelMultiAdapterLinear.

• MllamaVisionSdpaAttention (extends nn.Module) - Scaled dot-product attention for vision encoder layers with QKV projection supporting adapter loading.

• MllamaVisionEncoderLayer (extends nn.Module) - Single encoder layer combining self-attention and MLP with optional gating (used in global transformer layers).

• MllamaVisionEncoder (extends nn.Module) - Stacked encoder layers forming the backbone of the vision processing pipeline, collects intermediate hidden states.

• MllamaPrecomputedAspectRatioEmbedding and MllamaPrecomputedPositionEmbedding (extend nn.Module) - Handle tile-based position embeddings for variable aspect ratio image inputs.

• MllamaVisionModel (extends nn.Module) - The full vision model that processes images through patch embedding, tile/position embeddings, the local transformer, and the global transformer. Outputs include intermediate layer features concatenated with final hidden states.

Text Components:

• MllamaTextCrossAttention (extends nn.Module) - Cross-attention layer enabling the text model to attend to vision encoder outputs, with separate Q projection for text and KV projections for vision features.

• FlashLlamaCrossLayer - A text layer variant that includes cross-attention in addition to self-attention and MLP.

• MllamaForConditionalGeneration (extends nn.Module) - The top-level model class that combines the vision model with a modified Flash Llama text model, managing the multi-modal projection and cross-attention state passing.

Utility Functions:

• apply_rotary_pos_emb, rotate_half, repeat_kv - Standard transformer utilities for rotary embeddings and key-value head repetition.
• _prepare_aspect_ratio_attention_mask - Creates attention masks for variable aspect ratio tile-based vision inputs.
• _prepare_cross_attention_mask - Prepares cross-attention masks between text and vision tokens.

Usage

This model is loaded by the MllamaCausalLM wrapper class when the LoRax server initializes a Meta Llama multimodal model (e.g., Llama 3.2 Vision). The vision forward pass is invoked during prefill when images are present, producing cross-attention states that persist across decode steps. The model supports LoRA adapter loading for both vision and text components.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/mllama.py
• Lines: 1-1096

Signature

class MllamaVisionMLP(nn.Module):
    def __init__(self, *, prefix, config, weights, layer_id, model_type):
        ...
    def forward(self, hidden_states: torch.Tensor, adapter_data: AdapterBatchData) -> torch.Tensor:
        ...

class MllamaVisionSdpaAttention(nn.Module):
    def __init__(self, *, prefix, config, weights, layer_id, model_type):
        ...
    def forward(self, hidden_state, attention_mask=None, adapter_data=None) -> torch.Tensor:
        ...

class MllamaVisionEncoderLayer(nn.Module):
    def __init__(self, *, prefix, config, weights, is_gated, layer_id, model_type):
        ...

class MllamaVisionEncoder(nn.Module):
    def __init__(self, *, prefix, config, weights, is_gated, num_layers, model_type):
        ...

class MllamaVisionModel(nn.Module):
    def __init__(self, *, prefix, config, weights):
        ...
    def forward(self, pixel_values, aspect_ratio_ids, attention_mask, adapter_data) -> torch.Tensor:
        ...

class MllamaTextCrossAttention(nn.Module):
    def __init__(self, *, prefix, config, weights, layer_idx):
        ...

class MllamaForConditionalGeneration(nn.Module):
    def __init__(self, prefix, config, weights):
        ...
    def vision_forward(self, pixel_values, aspect_ratio_ids, aspect_ratio_mask, adapter_data):
        ...
    def forward(self, input_ids, position_ids, cu_seqlen_prefill, kv_cache, block_tables,
                slots, seqlen, max_s, prefill_cache_indices, lm_head_indices=None,
                cross_attention_states=None, adapter_data=None, image_indices=None):
        ...

Import

from lorax_server.models.custom_modeling.mllama import MllamaForConditionalGeneration

I/O Contract

Inputs

Name	Type	Required	Description
pixel_values	torch.Tensor	No	Image pixel values with shape (batch, media, tiles, channels, height, width)
aspect_ratio_ids	torch.Tensor	No	Aspect ratio identifiers for tile-based processing
aspect_ratio_mask	torch.Tensor	No	Mask for valid aspect ratio tiles
input_ids	torch.Tensor	Yes	Token IDs for the text input
position_ids	torch.Tensor	Yes	Position indices for each token
cu_seqlen_prefill	Optional[torch.Tensor]	Yes	Cumulative sequence lengths for prefill
kv_cache	List[Tuple[torch.Tensor, torch.Tensor]]	Yes	Key-value cache for each layer
block_tables	torch.Tensor	Yes	Block table mapping for paged attention
slots	torch.Tensor	Yes	Slot indices for KV cache storage
seqlen	Seqlen	Yes	Sequence length metadata wrapper
cross_attention_states	Optional[torch.Tensor]	No	Vision encoder outputs for cross-attention layers
adapter_data	Optional[AdapterBatchData]	No	LoRA adapter weights for the current batch
image_indices	List[int]	No	Mapping of cross-attention states to batch requests

Outputs

Name	Type	Description
logits	torch.Tensor	Next-token logits over the vocabulary
speculative_logits	Optional[torch.Tensor]	Speculative decoding logits (None if not using speculation)

Usage Examples

# Internal LoRax server usage
from lorax_server.models.custom_modeling.mllama import MllamaForConditionalGeneration

# Instantiated by MllamaCausalLM during model loading
# model = MllamaForConditionalGeneration(prefix="", config=config, weights=weights)

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime