Implementation:Predibase Lorax MPT Modeling

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

Overview

Provides a tensor-parallel implementation of the MPT (MosaicML Pretrained Transformer) causal language model with support for multiple attention backends including Flash Attention, Triton, and PyTorch.

Description

This module implements the MPT architecture, a GPT-style transformer with ALiBi positional encoding and configurable attention implementations. MPT supports both multi-head attention and multi-query attention variants.

Key classes:

• MultiheadAttention -- Multi-head self-attention with configurable backend (flash, triton, or torch). Uses TensorParallelColumnLinear for the fused QKV projection (Wqkv) and TensorParallelRowLinear for the output projection. Supports optional QKV clipping and QK layer normalization.

• MultiQueryAttention -- Multi-query attention variant where keys and values share a single head while queries have multiple heads. Same backend support as MultiheadAttention.

• MPTMLP -- Two-layer feed-forward network with configurable activation. Uses TensorParallelColumnLinear for up-projection and TensorParallelRowLinear for down-projection.

• MPTBlock -- Single transformer block with configurable normalization (supports LPLayerNorm and RMSNorm), attention (multi-head or multi-query), and MLP. Implements pre-norm architecture with optional residual scaling.

• MPTModel -- The full transformer backbone. Manages token embeddings (TensorParallelEmbedding), the block stack, final layer norm, and ALiBi attention bias construction. Requires ALiBi to be enabled (no learned positional embeddings).

• MPTForCausalLM -- Top-level causal LM wrapper. Contains MPTModel and an LM head (TensorParallelHead) with optional logit scaling. Returns CausalLMOutputWithPast.

Attention backends: scaled_multihead_dot_product_attention (PyTorch), flash_attn_fn (Flash Attention), triton_flash_attn_fn (Triton).

Norm classes: LPLayerNorm (low-precision LayerNorm), RMSNorm, LPRMSNorm.

Usage

Used internally by the LoRAX server when serving MPT-based models (e.g., mosaicml/mpt-7b, mpt-30b). Loaded via the model registry.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/mpt_modeling.py
• Lines: 1-1014

Signature

class MPTForCausalLM(MPTPreTrainedModel):
    def __init__(self, config, weights):
        ...

    def forward(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,
        attention_mask: Optional[torch.ByteTensor] = None,
        prefix_mask: Optional[torch.ByteTensor] = None,
        sequence_id: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        use_cache: Optional[bool] = None,
    ):

Import

from lorax_server.models.custom_modeling.mpt_modeling import MPTForCausalLM

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	torch.LongTensor	Yes	Input token IDs of shape (batch_size, sequence_length)
past_key_values	Optional[List[Tuple[torch.FloatTensor]]]	No	Cached key/value states for autoregressive decoding
attention_mask	Optional[torch.ByteTensor]	No	Attention mask for padding tokens
prefix_mask	Optional[torch.ByteTensor]	No	Prefix mask for prefix-LM mode
sequence_id	Optional[torch.LongTensor]	No	Sequence IDs for document-level attention masking
labels	Optional[torch.LongTensor]	No	Labels for language modeling loss
use_cache	Optional[bool]	No	Whether to return KV cache for autoregressive decoding

Outputs

Name	Type	Description
loss	Optional[torch.Tensor]	Language modeling loss (when labels provided)
logits	torch.Tensor	Prediction logits of shape (batch_size, sequence_length, vocab_size)
past_key_values	Tuple	Cached KV states for subsequent decoding
hidden_states	Optional[Tuple[torch.Tensor]]	Hidden states from all layers (when requested)
attentions	Optional[Tuple[torch.Tensor]]	Attention weights (when requested)

Usage Examples

# Internal usage within LoRAX server
from lorax_server.models.custom_modeling.mpt_modeling import MPTForCausalLM
# Instantiated by model registry with MPT config and pre-loaded weights

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime