Implementation:Predibase Lorax LLaVA NeXT Model

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

Overview

Implements the LLaVA-NeXT (Large Language and Vision Assistant - Next) multimodal model architecture, combining a vision tower, a multi-modal projector, and a text language model for vision-language conditional generation within the LoRax inference server.

Description

This module defines the PyTorch model components for LLaVA-NeXT, a vision-language model that processes images at any resolution using adaptive grid-based patching.

Key classes:

• LlavaNextMultiModalProjector (extends nn.Module) - A two-layer MLP that projects image features from the vision encoder's hidden dimension into the text model's hidden dimension. Uses TensorParallelColumnLinear and TensorParallelRowLinear for distributed inference.

• LlavaNextForConditionalGeneration (extends nn.Module) - The main model class that composes a vision tower, the multi-modal projector, and a text language model. During forward pass, it processes images through the vision tower, projects them via the MLP, and merges image features into the text embeddings at image token positions.

Key functions:

• get_anyres_image_grid_shape - Calculates the image patch grid shape for any-resolution image preprocessing by selecting the best resolution from configured grid pinpoints.

• unpad_image - Removes padding from resized images to maintain original aspect ratios, used during the spatial unpadding merge strategy.

Forward pass flow:

1. Input images are processed through the vision tower (CLIP-based)
2. Vision features are selected based on vision_feature_select_strategy ("default" removes CLS token, "full" keeps all)
3. Features are projected through the multi-modal projector
4. For multi-tile images, features are spatially rearranged, unpadded, and concatenated with newline embeddings
5. Image features are merged into text embeddings at image_token_index positions
6. The combined embeddings are passed through the text model with KV-cache support

Usage

This model class is instantiated by the VlmCausalLM wrapper when the LoRax server loads a LLaVA-NeXT model. It is passed as the model_class parameter and is used for both prefill and decode phases during inference. The vision tower is only invoked during prefill when pixel values are present.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/llava_next.py
• Lines: 1-273

Signature

def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
    ...

def unpad_image(tensor, original_size):
    ...

class LlavaNextMultiModalProjector(nn.Module):
    def __init__(self, prefix, config, weights):
        ...
    def forward(self, image_features):
        ...

class LlavaNextForConditionalGeneration(nn.Module):
    def __init__(self, prefix, config, weights):
        ...
    def _merge_input_ids_with_image_features(
        self,
        input_ids: torch.Tensor,
        inputs_embeds: torch.Tensor,
        image_features: torch.Tensor,
    ):
        ...
    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlen_prefill: Optional[torch.Tensor],
        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
        block_tables: torch.Tensor,
        slots: torch.Tensor,
        seqlen: Seqlen,
        max_s: int,
        prefill_cache_indices: Optional[torch.Tensor],
        lm_head_indices: Optional[torch.Tensor] = None,
        pixel_values: torch.FloatTensor = None,
        pixel_attention_mask=None,
        image_sizes: Optional[torch.LongTensor] = None,
        adapter_data: Optional[AdapterBatchData] = None,
        skip_lm_head: bool = False,
    ):
        ...

Import

from lorax_server.models.custom_modeling.llava_next import LlavaNextForConditionalGeneration

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	torch.Tensor	Yes	Token IDs for the text input sequence
position_ids	torch.Tensor	Yes	Position indices for each token
cu_seqlen_prefill	Optional[torch.Tensor]	Yes	Cumulative sequence lengths for prefill (None during decode)
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
max_s	int	Yes	Maximum sequence length in the batch
prefill_cache_indices	Optional[torch.Tensor]	Yes	Cache indices for prefix caching
pixel_values	torch.FloatTensor	No	Image pixel values with shape (num_images, num_patches, C, H, W)
image_sizes	Optional[torch.LongTensor]	No	Original sizes of each image (height, width)
adapter_data	Optional[AdapterBatchData]	No	LoRA adapter weights for the current batch

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.llava_next import LlavaNextForConditionalGeneration

# Instantiated by VlmCausalLM during model loading
# model = LlavaNextForConditionalGeneration(prefix="", config=config, weights=weights)

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime