Implementation:Predibase Lorax GPTQ Exllama V1
| Knowledge Sources | |
|---|---|
| Domains | Quantization, Inference |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
Provides a 4-bit quantized linear layer using the ExLlama v1 CUDA kernels for efficient GPTQ inference with optional activation-order (act-order) support.
Description
This module wraps the exllama_kernels C++ extension to perform 4-bit quantized matrix multiplication on CUDA devices. It contains the following components:
ext_make_q4: A helper function that constructs a Q4Matrix handle from quantized weights, zero points, scales, and an optional group index tensor. A dummy none_tensor on the meta device is passed when g_idx is None, since the C++ extension cannot accept Python None.
ext_q4_matmul: Performs matrix multiplication x @ q4 by calling the q4_matmul kernel, reshaping input to 2D and producing output of shape (batch, q4_width).
create_exllama_buffers: Allocates temporary GPU buffers (temp_state and temp_dq) needed by the ExLlama kernels. The temp_state buffer is used to reorder activations in the act-order case, and temp_dq is used for weight dequantization with cuBLAS during prefill. Also configures tuning parameters via set_tuning_params.
Ex4bitLinear: An nn.Module that represents a 4-bit quantized linear layer. During initialization, it detects whether the model uses activation ordering by checking if g_idx is non-trivial. It tracks global MAX_DQ and MAX_INNER values used to size shared temporary buffers. The forward method calls ext_q4_matmul and optionally adds bias.
Usage
This layer is used when loading GPTQ-quantized models with ExLlama v1 acceleration enabled. The create_exllama_buffers function must be called after all layers are initialized to allocate shared temporary GPU memory before inference begins.
Code Reference
Source Location
- Repository: Predibase_Lorax
- File: server/lorax_server/layers/gptq/exllama.py
- Lines: 1-125
Signature
class Ex4bitLinear(torch.nn.Module):
def __init__(self, qweight, qzeros, scales, g_idx, bias, bits, groupsize):
Import
from lorax_server.layers.gptq.exllama import Ex4bitLinear, create_exllama_buffers
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| qweight | torch.Tensor (int32) | Yes | Packed 4-bit quantized weight matrix on CUDA |
| qzeros | torch.Tensor (int32) | Yes | Packed quantized zero points |
| scales | torch.Tensor (float16) | Yes | Per-group scale factors |
| g_idx | torch.Tensor (int32) or None | No | Group index for activation ordering; None or trivial indices indicate no act-order |
| bias | torch.Tensor or None | No | Optional bias vector |
| bits | int | Yes | Must be 4 (only 4-bit quantization is supported) |
| groupsize | int | Yes | Number of input features per quantization group |
Outputs
| Name | Type | Description |
|---|---|---|
| output | torch.Tensor (float16) | Result of the quantized linear transformation |
Usage Examples
# Used internally by model layers
from lorax_server.layers.gptq.exllama import Ex4bitLinear, set_device, create_exllama_buffers
set_device(torch.device("cuda:0"))
layer = Ex4bitLinear(qweight, qzeros, scales, g_idx, bias, bits=4, groupsize=128)
create_exllama_buffers(max_total_tokens=2048)
output = layer(input_tensor)