Implementation:Predibase Lorax Flash SantaCoder Modeling

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

Overview

Provides a Flash Attention-based, tensor-parallel implementation of the SantaCoder causal language model with multi-query attention and paged KV cache for high-throughput code generation inference within the LoRAX serving framework.

Description

This module implements the SantaCoder architecture using Flash Attention and paged attention. SantaCoder is a GPT-style code generation model that uses multi-query attention (MQA), where all query heads share a single key and value head, enabling more efficient KV caching. The model uses learned positional embeddings (both token and position embeddings).

Key classes:

• FlashMQAttention -- Flash multi-query attention implementation. Uses a custom load_multi_mqa function to load the fused QKV weights where the query projection is sharded across tensor-parallel ranks but the key/value projections (single head each) are replicated. The kv_head_mapping maps all query heads to the single KV head (index 0). Uses flash_attn for prefill and paged_attention for decode.

• MLP -- Two-layer feed-forward network with configurable activation (GeLU variants). Uses load_col for the column-parallel up-projection (c_fc) and load_row for the row-parallel down-projection (c_proj).

• Block -- Single transformer block with FastLayerNorm, multi-query attention, and MLP. Implements pre-norm architecture. Attention and MLP are computed sequentially with separate residual connections.

• FlashSantacoderModel -- Full transformer backbone. Combines token embeddings (wte) and positional embeddings (wpe) via TensorParallelEmbedding (with reduce=False, performing an all-reduce after summation). Runs the block stack and applies a final FastLayerNorm.

• FlashSantacoderForCausalLM -- Top-level causal LM wrapper with FlashSantacoderModel and an LM head (TensorParallelHead tied to wte). Accepts paged KV cache inputs and returns logits.

Helper functions: load_multi_mqa handles the complex weight loading for multi-query attention, supporting both standard and GPTQ-quantized weight formats, as well as transposed weight layouts. _load_multi_mqa_gptq and _load_multi_mqa are the quantized and non-quantized variants respectively.

Usage

Used internally by the LoRAX server when serving SantaCoder and StarCoder-based models. Loaded via the model registry.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/flash_santacoder_modeling.py
• Lines: 1-445

Signature

class FlashSantacoderForCausalLM(nn.Module):
    def __init__(self, prefix: str, config, weights):
        ...

    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] = None,
        lm_head_indices: Optional[torch.Tensor] = None,
        skip_lm_head: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

Import

from lorax_server.models.custom_modeling.flash_santacoder_modeling import FlashSantacoderForCausalLM

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	torch.Tensor	Yes	Input token IDs (flattened across batch)
position_ids	torch.Tensor	Yes	Position IDs for learned positional embeddings
cu_seqlen_prefill	Optional[torch.Tensor]	No	Cumulative sequence lengths for prefill (Flash Attention format). None during decode.
kv_cache	List[Tuple[torch.Tensor, torch.Tensor]]	Yes	Paged KV cache tensors, one pair per layer
block_tables	torch.Tensor	Yes	Block table mapping for paged attention
slots	torch.Tensor	Yes	Slot indices for KV cache insertion
seqlen	Seqlen	Yes	Sequence length information for paged attention
max_s	int	Yes	Maximum sequence length in the batch
lm_head_indices	Optional[torch.Tensor]	No	Indices to select specific positions for LM head
skip_lm_head	bool	No	If True, return hidden states instead of logits

Outputs

Name	Type	Description
logits	torch.Tensor	Prediction logits (or hidden states if skip_lm_head=True)
speculative_logits	Optional[torch.Tensor]	Always None for this model

Usage Examples

# Internal usage within LoRAX server
from lorax_server.models.custom_modeling.flash_santacoder_modeling import FlashSantacoderForCausalLM
# Instantiated by model registry with SantaCoder config and pre-loaded weights
# Uses multi-query attention for efficient KV caching

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime