Implementation:Predibase Lorax Linear Layer Factory

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

Overview

Provides the FastLinear base linear layer and the get_linear factory function that dispatches to the appropriate quantized or unquantized linear layer implementation based on the quantization method.

Description

This module serves as the central factory for creating linear layers in LoRAX, with support for multiple quantization backends:

FastLinear: A simple nn.Module wrapping F.linear with weight and optional bias as non-trainable parameters. It provides a load classmethod for loading weights from checkpoints. This is the default layer used when no quantization is specified.

FastLinearROCm: A ROCm-optimized variant of FastLinear that uses custom CUDA kernels from vLLM (_custom_C.LLMM1) for single-token inference when the batch dimension is 1 and the weight dimensions match specific patterns (e.g., k=8192 with m=1280 or m=7168). Falls back to F.linear for other cases.

get_linear: The factory function that creates the appropriate linear layer based on the quantize parameter:

• None: Returns FastLinear
• fp8 / fp8-kv: Returns Fp8Linear from lorax_server.layers.fp8
• bitsandbytes: Returns Linear8bitLt with 8-bit quantization
• bitsandbytes-nf4: Returns Linear4bit with NF4 quantization
• bitsandbytes-fp4: Returns Linear4bit with FP4 quantization
• eetq: Returns EETQLinear for W8A16 quantization
• gptq: Unpacks a 7-tuple (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama) and returns either exllamav2.QuantLinear (when use_exllama is True) or quant_linear.QuantLinear
• awq: Returns AWQLinear from the AWQ backend
• hqq-*: Returns an HQQ linear layer

The function also handles fan_in_fan_out transposition for Conv1D-to-Linear replacement (used by some LoRA configurations).

Usage

This is the core factory used by all tensor-parallel layer classes (TensorParallelColumnLinear, TensorParallelRowLinear, TensorParallelHead) to instantiate the correct linear layer for the model's quantization configuration.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/layers/linear.py
• Lines: 1-181

Signature

class FastLinear(torch.nn.Module):
    def __init__(self, weight, bias) -> None:

class FastLinearROCm(torch.nn.Module):
    def __init__(self, weight, bias) -> None:

def get_linear(weight, bias, quantize, fan_in_fan_out=False,
               weight_scale=None, input_scale=None):

Import

from lorax_server.layers.linear import FastLinear, get_linear

I/O Contract

Inputs (get_linear)

Name	Type	Required	Description
weight	torch.Tensor or tuple	Yes	Weight tensor, or a tuple of quantized weight components for GPTQ/AWQ
bias	torch.Tensor or None	No	Optional bias tensor
quantize	str or None	Yes	Quantization method name (None, "gptq", "awq", "bitsandbytes", "eetq", "fp8", "hqq-*")
fan_in_fan_out	bool	No	If True, transpose weight for Conv1D-to-Linear replacement (default False)
weight_scale	torch.Tensor or None	No	Weight scale for FP8 quantization
input_scale	torch.Tensor or None	No	Input scale for FP8 quantization

Outputs

Name	Type	Description
linear	torch.nn.Module	An instantiated linear layer of the appropriate type for the quantization method

Usage Examples

# Used internally by tensor-parallel layer classes
from lorax_server.layers.linear import FastLinear, get_linear

# Unquantized
linear = get_linear(weight, bias=None, quantize=None)

# GPTQ quantized
linear = get_linear((qweight, qzeros, scales, g_idx, 4, 128, True), bias=None, quantize="gptq")

# 8-bit bitsandbytes
linear = get_linear(weight, bias=None, quantize="bitsandbytes")

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime