Implementation:Predibase Lorax Flash Solar Modeling

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

Overview

Provides a Flash Attention-based, adapter-aware, tensor-parallel implementation of the Solar causal language model with cross-layer residual connections (BSKCN) and paged KV cache for high-throughput inference within the LoRAX serving framework.

Description

This module implements the Solar architecture, which extends a LLaMA-like design with a distinctive cross-layer residual connection mechanism called BSKCN (Block-Skip-Connection-Network). Solar uses grouped-query attention with rotary positional embeddings and supports LoRA adapter injection. Flash Attention v2 is required.

Key classes:

• SolarConfig -- Configuration class extending PretrainedConfig with Solar-specific parameters: bskcn_1, bskcn_2, bskcn_3, bskcn_4 (layer index lists for cross-layer residual connections), and bskcn_tv (interpolation weights for training/inference). Also includes sliding_window, rope_scaling, and rope_theta for attention configuration.

• SolarRMSNorm -- RMS normalization with fused residual addition via dropout_layer_norm kernel for sequences up to 8192 hidden dimensions. Falls back to manual computation for larger dimensions.

• SolarAttention -- Grouped-query attention with rotary positional embeddings. Uses TensorParallelMultiAdapterLinear for the QKV projection (supporting LoRA on Q, K, V independently) and TensorParallelAdapterRowLinear for the output projection (supporting LoRA on O). Computes RoPE via PositionRotaryEmbedding. Supports FP8 KV cache quantization. Delegates to flash_attn for prefill and paged_attention for decode.

• SolarMLP -- Gated feed-forward network (SwiGLU-style) with gate and up projections via TensorParallelMultiAdapterLinear and down projection via TensorParallelAdapterRowLinear. Supports LoRA adapters on gate_proj, up_proj, and down_proj.

• SolarLayer -- Single transformer layer combining SolarRMSNorm, SolarAttention, and SolarMLP with residual connections. Uses pre-norm architecture with fused residual addition in the layer norm kernels.

• SolarModel -- Full transformer backbone. Implements the BSKCN cross-layer residual mechanism: at designated layer indices (bskcn_1, bskcn_2), hidden states are saved; at other designated indices (bskcn_3, bskcn_4), saved states are interpolated with current states using bskcn_tv weights. This creates skip connections across non-adjacent layers.

• FlashSolarForCausalLM -- Top-level causal LM wrapper. Contains token embeddings (TensorParallelEmbedding), SolarModel, and an LM head (MultiAdapterHead wrapping TensorParallelHead). Supports sliding window attention via max_past. Accepts paged KV cache inputs and returns logits plus optional speculative logits from the adapter head.

Usage

Used internally by the LoRAX server when serving Solar-based models (e.g., upstage/SOLAR-10.7B). Loaded via the model registry.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/models/custom_modeling/flash_solar_modeling.py
• Lines: 1-699

Signature

class FlashSolarForCausalLM(torch.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,
        adapter_data: AdapterBatchData,
        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_solar_modeling import FlashSolarForCausalLM

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
adapter_data	AdapterBatchData	Yes	LoRA adapter batch data for dynamic adapter application
prefill_cache_indices	Optional[torch.Tensor]	No	Indices for selective KV cache population during prefill
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]	Optional speculative decoding logits from the multi-adapter head

Usage Examples

# Internal usage within LoRAX server
from lorax_server.models.custom_modeling.flash_solar_modeling import FlashSolarForCausalLM
# Instantiated by model registry with SolarConfig and pre-loaded weights
# Supports LoRA adapters on attention (Q, K, V, O) and MLP (gate, up, down) projections
# Uses BSKCN cross-layer residual connections for improved performance

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime