Implementation:Predibase Lorax Bloom Modeling

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

Overview

Provides a tensor-parallel, adapter-aware implementation of the BLOOM causal language model for optimized inference within the LoRAX serving framework.

Description

This module implements the BLOOM (BigScience Large Open-science Open-access Multilingual) language model architecture with tensor parallelism and LoRA adapter support. It uses ALiBi (Attention with Linear Biases) positional encoding instead of traditional positional embeddings.

Key classes:

• BloomAttention -- Multi-head attention with ALiBi positional biases and optional fused CUDA kernels. Supports adapter injection on the QKV projection (query_key_value) via TensorParallelMultiAdapterLinear and on the dense output via TensorParallelAdapterRowLinear. When custom kernels are enabled, the attention is computed via fused_bloom_attention_cuda for improved performance.

• BloomMLP -- Two-layer feed-forward network with GeLU activation, using TensorParallelMultiAdapterLinear for the up-projection and TensorParallelAdapterRowLinear for the down-projection to support LoRA adapters.

• BloomBlock -- A single transformer block combining layer normalization, self-attention with residual connection (via dropout_add), and MLP with residual connection. Implements pre-norm architecture.

• BloomModel -- The full transformer backbone. Manages word embeddings (TensorParallelEmbedding), word embeddings layer norm, the stack of BloomBlock layers, and a final layer norm. Constructs ALiBi attention biases and causal masks.

• BloomForCausalLM -- Top-level causal language model wrapper. Contains BloomModel and an LM head (TensorParallelHead with adapter support). The forward method returns CausalLMOutputWithCrossAttentions with logits and optional KV cache.

Helper functions: _make_causal_mask, _expand_mask, build_alibi_tensor, dropout_add, _split_heads, _merge_heads.

Usage

Used internally by the LoRAX server when serving BLOOM-based models (e.g., bigscience/bloom-560m, bloom-7b1). Loaded via the model registry.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/bloom_modeling.py
• Lines: 1-913

Signature

class BloomForCausalLM(BloomPreTrainedModel):
    def __init__(self, config, weights):
        ...

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        adapter_data: Optional[AdapterBatchData] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:

Import

from lorax_server.models.custom_modeling.bloom_modeling import BloomForCausalLM

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	Optional[torch.LongTensor]	No	Input token IDs of shape (batch_size, sequence_length)
past_key_values	Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]	No	Cached key/value states for autoregressive decoding
attention_mask	Optional[torch.Tensor]	No	Attention mask of shape (batch_size, sequence_length)
head_mask	Optional[torch.Tensor]	No	Per-head attention mask
inputs_embeds	Optional[torch.Tensor]	No	Pre-computed input embeddings (alternative to input_ids)
labels	Optional[torch.Tensor]	No	Labels for language modeling loss computation
use_cache	Optional[bool]	No	Whether to return KV cache for autoregressive decoding
adapter_data	Optional[AdapterBatchData]	No	LoRA adapter batch data for dynamic adapter application

Outputs

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

Usage Examples

# Internal usage within LoRAX server
from lorax_server.models.custom_modeling.bloom_modeling import BloomForCausalLM
# Instantiated by model registry with BloomConfig and pre-loaded weights

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime