Implementation:Predibase Lorax MLlama Wrapper

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

Overview

Provides the MLlama (Meta Llama Vision) inference wrapper that extends VlmCausalLM with MLlama-specific batch handling, cross-attention state management, and full LoRA adapter support for both text and vision transformer layers.

Description

This module implements the server-level wrapper for Meta's Llama multimodal models, handling the unique cross-attention architecture where vision encoder outputs persist across decode steps.

Key classes:

• MllamaCausalLMBatch (extends VlmCausalLMBatch, dataclass) - A specialized batch type that adds:
  • image_indices - Tracks which batch requests have associated images for cross-attention.
  • aspect_ratio_ids and aspect_ratio_mask - Metadata for variable aspect ratio image tiles.
  • cross_attention_states - Vision encoder output tensor that persists across decode steps.
  • batch_tokenized_inputs - Processes <|image|> tokens and runs per-image processing through the MLlama image processor independently (unlike other VLMs that batch-process images).
  • concatenate - Merges cross-attention states from multiple batches with proper index offset tracking.
  • filter - Filters cross-attention states when requests are removed, maintaining correct index mappings.
  • from_pb - Clamps out-of-bounds <|image|> token IDs to prevent logit processor errors.

• MllamaCausalLM (extends VlmCausalLM) - The model wrapper that:
  • Enables adapter loading (supports_adapter_loading = True) for both text and vision layers.
  • Defines adapter layers spanning text (Q, K, V, O, GATE, UP, DOWN projections plus LM_HEAD) and vision (VISION_GLOBAL_TRANSFORMER and VISION_TRANSFORMER with Q, K, V, O, FC1, FC2).
  • Implements adapter_target_to_layer to map adapter weight names to model parameters, with special handling for FlashLlamaCrossLayer cross-attention layers vs standard self-attention layers.
  • Always uses adapter_prefill_state(prefill=True) to force SGMV kernels (instead of BGMV) because cross-attention layers require this during decode.
  • Overrides forward to invoke vision_forward on pixel values during prefill, store the resulting cross-attention states, and pass them to the text model on every subsequent step.

Usage

MllamaCausalLM is instantiated by the LoRax model registry when loading Meta Llama Vision models. It manages the two-phase inference: vision encoding during prefill (producing cross-attention states) and text generation with cross-attention on every decode step.

Code Reference

Source Location

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

Signature

@dataclass
class MllamaCausalLMBatch(VlmCausalLMBatch):
    image_indices: List[int] = 42
    aspect_ratio_ids: Optional[torch.Tensor] = None
    aspect_ratio_mask: Optional[torch.Tensor] = None
    cross_attention_states: Optional[torch.Tensor] = None

    @classmethod
    def concatenate(cls, batches):
        ...
    def filter(self, request_ids: List[int]):
        ...
    @classmethod
    def batch_tokenized_inputs(cls, requests, tokenizer, processor, config):
        ...
    @classmethod
    def from_pb(cls, pb, tokenizer, tokenizers, processor, config, dtype, device) -> "VlmCausalLMBatch":
        ...

class MllamaCausalLM(VlmCausalLM):
    @property
    def supports_adapter_loading(self) -> bool:
        ...
    @property
    def adapter_layers(self) -> List[str]:
        ...
    @property
    def default_traced_adapter_layers(self) -> List[str]:
        ...
    def get_num_layers_for_type(self, layer_type: str) -> int:
        ...
    def adapter_target_to_layer(self) -> Dict[str, Tuple[str, torch.Tensor]]:
        ...
    def adapter_prefill_state(self, prefill: bool) -> bool:
        ...
    def forward(self, batch: VlmCausalLMBatch, adapter_data=None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        ...

Import

from lorax_server.models.mllama import MllamaCausalLM, MllamaCausalLMBatch

I/O Contract

Inputs

Name	Type	Required	Description
batch	VlmCausalLMBatch	Yes	Batch containing text tokens, pixel values, aspect ratio metadata, and cross-attention states
adapter_data	Optional[Dict[str, torch.Tensor]]	No	Adapter weight data for text and vision layers

Outputs

Name	Type	Description
logits	torch.Tensor	Next-token logits over the vocabulary
speculative_logits	Optional[torch.Tensor]	Speculative decoding logits

Usage Examples

# Internal LoRax server usage
from lorax_server.models.mllama import MllamaCausalLM

# Instantiated by model registry for Meta Llama Vision models
# mllama = MllamaCausalLM(
#     model_id="meta-llama/Llama-3.2-11B-Vision",
#     model_class=MllamaForConditionalGeneration,
#     adapter_id="",
#     adapter_source="hub",
# )

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime