Implementation:Predibase Lorax Tensor Parallel Layers

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

Overview

Implements tensor-parallel linear, embedding, and head layers that shard model parameters across multiple GPU processes using PyTorch distributed all-gather and all-reduce operations.

Description

This module provides the tensor-parallel infrastructure used by all model implementations in LoRAX for multi-GPU inference:

SuperLayer: A thin wrapper nn.Module that holds a linear layer and delegates forward to it. Serves as the base class for all tensor-parallel layers.

TensorParallelHead: Extends SuperLayer for the language model head (final projection to vocabulary). The load staticmethod loads weights sharded along dimension 0 across processes, falling back to loading the full tensor if the vocabulary size is not evenly divisible. During forward, when should_gather is True, it performs an optimized all-gather: for single-token batches with FastLinear, it fuses the matmul and all-gather into tensor operations to minimize memory allocations. For multi-token batches, it performs a transposed matmul followed by all_gather_into_tensor. For non-FastLinear backends, it uses standard torch.distributed.all_gather with concatenation.

TensorParallelColumnLinear: Shards linear layer weights along the output dimension (columns). Provides multiple classmethods:

• load: Loads a single weight tensor sharded along dim 0.
• load_multi: Loads and concatenates multiple weight tensors (e.g., Q, K, V projections) along a specified dimension.
• load_qkv: Loads packed QKV weights for models that fuse these projections.
• load_gate_up: Loads packed gate/up projections for gated MLPs.

TensorParallelRowLinear: Shards linear layer weights along the input dimension (rows). The forward method performs an all_reduce when the process group has more than one member and reduce is True. Only rank 0 loads the bias.

TensorParallelEmbedding: Shards embedding tables across processes by partitioning the vocabulary range. Each process holds a block of embeddings from min_id to max_id. Out-of-range indices are mapped to a padding entry (zero vector) using torch.where. The forward method applies the embedding lookup and optionally performs an all_reduce to combine partial results.

Usage

These classes are used by every model implementation in LoRAX. TensorParallelColumnLinear is used for attention projections and MLP up-projections, TensorParallelRowLinear for attention output and MLP down-projections, TensorParallelHead for the final LM head, and TensorParallelEmbedding for token embeddings.

Code Reference

Source Location

• Repository: Predibase_Lorax
• File: server/lorax_server/layers/tensor_parallel.py
• Lines: 1-192

Signature

class SuperLayer(torch.nn.Module):
    def __init__(self, linear):

class TensorParallelHead(SuperLayer):
    def __init__(self, linear, process_group, should_gather: bool):

class TensorParallelColumnLinear(SuperLayer):
    # classmethods: load, load_multi, load_qkv, load_gate_up

class TensorParallelRowLinear(SuperLayer):
    def __init__(self, linear, process_group):

class TensorParallelEmbedding(torch.nn.Module):
    def __init__(self, prefix: str, weights, reduce=True):

Import

from lorax_server.layers.tensor_parallel import (
    TensorParallelHead,
    TensorParallelColumnLinear,
    TensorParallelRowLinear,
    TensorParallelEmbedding,
)

I/O Contract

Inputs (TensorParallelColumnLinear.load)

Name	Type	Required	Description
config	object	Yes	Model configuration with quantize attribute
prefix	str	Yes	Weight name prefix in the checkpoint
weights	Weights	Yes	Weight loader with process_group information
bias	bool	Yes	Whether to load bias

Inputs (TensorParallelEmbedding.__init__)

Name	Type	Required	Description
prefix	str	Yes	Weight name prefix for the embedding table
weights	Weights	Yes	Weight loader with process_group and sharding info
reduce	bool	No	Whether to all-reduce after embedding lookup (default True)

Outputs

Name	Type	Description
output	torch.Tensor	Result of the tensor-parallel linear or embedding operation, gathered or reduced across processes

Usage Examples

# Used internally by model implementations
from lorax_server.layers.tensor_parallel import TensorParallelColumnLinear, TensorParallelRowLinear

qkv_proj = TensorParallelColumnLinear.load_multi(
    config, prefixes=["q_proj", "k_proj", "v_proj"], weights=weights, bias=False, dim=0
)
o_proj = TensorParallelRowLinear.load(config, prefix="o_proj", weights=weights, bias=False)

Related Pages

• Environment:Predibase_Lorax_CUDA_GPU_Runtime