Implementation:Predibase Lorax Flash RW 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 RefinedWeb (RW/Falcon) causal language model with paged KV cache for high-throughput inference within the LoRAX serving framework.

Description

This module implements the RefinedWeb/Falcon architecture using Flash Attention and paged attention for efficient serving. It supports both the original Falcon architecture (multi-query attention) and the newer "large" decoder architecture (grouped-query attention). The model uses rotary positional embeddings (RoPE).

Key classes:

• RWConfig -- Custom configuration class for RefinedWeb models. Supports configurable number of KV heads (multi-query or multi-head), parallel attention (attention and MLP computed in parallel), and the newer decoder architecture variant. ALiBi is explicitly not supported; only rotary embeddings are used.

• FlashRWAttention -- Flash attention implementation for the original Falcon architecture. Uses a fused query_key_value projection via TensorParallelColumnLinear. Supports multi-query attention (single KV head shared across all query heads) with kv_head_mapping. Uses PositionRotaryEmbedding for RoPE, and delegates to flash_attn for prefill and paged_attention for decode.

• FlashRWLargeAttention -- Flash attention for the newer "large" Falcon decoder architecture. Supports grouped-query attention with configurable number of KV groups. Uses separate group-level QKV projections.

• FlashMLP -- Two-layer feed-forward network with GeLU activation, using TensorParallelColumnLinear for up-projection and TensorParallelRowLinear for down-projection.

• FlashRWLayer -- Transformer layer for the original architecture. Supports parallel attention and MLP (when config.parallel_attn=True), where attention and MLP outputs are summed with the residual in a single step. Uses FastLayerNorm.

• FlashRWLargeLayer -- Transformer layer for the newer large decoder architecture. Uses separate layer norms for self-attention and MLP with sequential residual additions.

• FlashRWModel -- Full transformer backbone. Dynamically selects between FlashRWLayer and FlashRWLargeLayer based on config.new_decoder_architecture. Manages word embeddings and computes rotary cos/sin once per forward pass.

• FlashRWForCausalLM -- Top-level causal LM wrapper with FlashRWModel and an LM head (TensorParallelHead). Accepts paged KV cache inputs and returns logits.

Usage

Used internally by the LoRAX server when serving Falcon/RefinedWeb-based models (e.g., tiiuae/falcon-7b, falcon-40b). Loaded via the model registry.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/flash_rw_modeling.py
• Lines: 1-614

Signature

class FlashRWForCausalLM(FlashRWPreTrainedModel):
    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_rw_modeling import FlashRWForCausalLM

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 rotary 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 (for next-token prediction)
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_rw_modeling import FlashRWForCausalLM
# Instantiated by model registry with RWConfig and pre-loaded weights
# Uses paged KV cache for efficient batch serving

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime