LoRAX, in depth: serving hundreds of LoRA adapters on one base model with SGMV, exchange scheduling and tiered weight caching

What LoRAX is and is not

LoRAX began as a fork of Hugging Face's text-generation-inference (TGI). It keeps that lineage's shape: a Rust launcher and router in front of Python model shards, flash and paged attention, tensor parallelism, token streaming and quantised base weights through --quantize (bitsandbytes, GPTQ and AWQ variants among the options). What it adds is a different unit of deployment. In a normal server the deployment is a model. In LoRAX the deployment is a base model, and the adapter is a request parameter.

Three features carry that idea. Dynamic adapter loading: an adapter is fetched the first time a request names it, from the Hugging Face Hub, a local directory, S3 or Predibase, with no restart. Heterogeneous continuous batching: requests for different adapters share the same decode batches. Adapter exchange scheduling: a scheduler decides which adapters are resident on the GPU and prefetches the next ones while the current ones run. Supported adapter types are LoRA and Medusa speculative-decoding heads. The hard constraint: one base model per deployment. Every adapter must have been trained against exactly that base, with the same tokenizer and architecture, or its outputs will be garbage even though loading succeeds.

Architecture: router, scheduler and three weight tiers

Follow one request. The client sends a prompt with adapter_id set. The router puts it in that adapter's queue. If the adapter is already active on the GPU, the request becomes eligible for the next batch immediately, exactly like a base-model request in TGI. If not, the adapter must be fetched from its source into host memory and then copied to the GPU, and the request waits in its queue while that happens. Other adapters' requests keep decoding; a cold start costs the tenant who triggered it, not everyone.

Memory on the GPU is split three ways: the base weights, a reservation for adapter weights controlled by --adapter-memory-fraction, and whatever is left for the paged KV cache. That split is the central capacity trade-off of the whole system, and the worked example below puts numbers on it.

Computing a batch that mixes adapters

A LoRA adapter replaces a frozen projection y = x W with y = x W + s * (x A) B, where A is d_in by r, B is r by d_out, r is the rank (often 8 to 64) and s is the scaling factor alpha divided by r. The base term is identical for every request in the batch, so it runs as one large GEMM over all tokens. Only the low-rank term differs by adapter.

The naive approach loops over adapters and launches two small matmuls each, which wastes the GPU on launch overhead and tiny shapes. LoRAX instead uses SGMV (segmented gather matrix-vector multiplication) kernels from the Punica project: tokens are arranged so that each adapter's tokens form a contiguous segment, and a single kernel launch gathers the right A and B for each segment. The reference computation looks like this:

import torch

def lora_mixed_batch(x, W, adapters, seg_starts, seg_adapter_ids, scale):
    """x: [tokens, d_in] sorted so each adapter's tokens are contiguous.
    seg_starts: segment boundaries, length n_segments + 1.
    seg_adapter_ids: which adapter owns each segment (None = base only).
    adapters[i] = (A_i [d_in, r], B_i [r, d_out])."""
    y = x @ W                                   # shared base GEMM, all tokens
    for k, aid in enumerate(seg_adapter_ids):   # SGMV fuses this loop into one launch
        if aid is None:
            continue
        lo, hi = seg_starts[k], seg_starts[k + 1]
        A, B = adapters[aid]
        y[lo:hi] += scale[aid] * ((x[lo:hi] @ A) @ B)   # shrink to r, expand back
    return y

Two properties follow. First, the extra cost of mixing adapters is small relative to the base GEMM, because the rank r is tiny next to the hidden size; the overhead grows with the number of distinct segments, not the number of tokens. Second, ranks do not need to match across adapters for correctness, but very different ranks and target modules make segments uneven, so benchmark with the adapter mix you will actually serve.

Adapter exchange scheduling

If every adapter that ever received a request stayed on the GPU, adapter memory would eventually crowd out the KV cache. The exchange scheduler bounds that. Two launcher flags define its behaviour: --max-active-adapters caps how many adapters can be active for request processing at once (the launcher default is 1024), and --adapter-cycle-time-s sets the interval between adapter exchanges (default 2 seconds).

Each cycle the scheduler looks at the per-adapter queues, decides which adapters should be active, offloads ones that have gone idle, and asynchronously prefetches ones with queued work. Within the active set, continuous batching admits requests as KV pages free up. The practical consequences for operators are about tails, not averages:

• A request for a resident adapter sees ordinary continuous-batching latency.
• A request for an adapter in host RAM pays a host-to-device copy, typically milliseconds for a small adapter.
• A request for a cold adapter pays a network fetch plus the copy. From the Hub or S3 this can be seconds, and it is the number your p99 time-to-first-token will show.
• When the number of hot adapters exceeds what fits, adapters churn: each cycle evicts something that will be needed again shortly. Churn shows up as rising TTFT with flat GPU utilisation.

Running it: launcher flags and requests

A minimal deployment pins the base model and gives the adapter reservation an explicit value, so capacity is a decision rather than a default:

model=mistralai/Mistral-7B-Instruct-v0.1
volume=$PWD/data

docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data \
  ghcr.io/predibase/lorax:main \
  --model-id $model \
  --adapter-memory-fraction 0.15 \
  --max-active-adapters 128 \
  --adapter-cycle-time-s 2

Pin the image to a release tag rather than :main once you are past experiments. A native request names the adapter and its source; adapter_source accepts hub (the default), local, s3 and pbase. A local adapter directory must contain adapter_config.json and the weights file, and S3 access reads the standard AWS credential variables.

import requests

r = requests.post("http://localhost:8080/generate", json={
    "inputs": "[INST] Classify this ticket: printer offline after update [/INST]",
    "parameters": {
        "max_new_tokens": 64,
        "adapter_id": "s3://adapters/support-triage/v7",
        "adapter_source": "s3",
    },
})
print(r.json()["generated_text"])

# OpenAI-compatible route: the model field selects the adapter ("" = base model)
from openai import OpenAI
client = OpenAI(api_key="EMPTY", base_url="http://localhost:8080/v1")
resp = client.chat.completions.create(
    model="acme/support-triage-lora",
    messages=[{"role": "user", "content": "printer offline after update"}],
    max_tokens=64,
)

Version adapters in their IDs (support-triage/v7, not support-triage/latest). The server caches by ID, so overwriting files behind an unchanged ID gives you a fleet where some replicas serve the old weights and some the new.

Worked example: how many adapters fit

Take Llama-3-8B (hidden size 4096, 32 layers, 8 key-value heads of dimension 128, so the K and V projections output 1024) with rank-16 LoRA on the four attention projections. Per layer, each projection adds r times (d_in + d_out) parameters:

On an 80 GB GPU the bf16 base weights take roughly 16 GB. A 10 percent adapter reservation of the remaining usable memory is in the region of 6 to 8 GB, which holds very roughly 230 to 290 such adapters, with the rest going to KV cache. Exact figures depend on how the server measures free memory after loading, so read the startup log rather than trusting the arithmetic. Now change one variable: the same adapters at rank 64 on all seven linear layers (attention plus MLP, where the MLP width is 14336) come to about 168 M parameters, roughly 335 MB each in fp16, and the same reservation holds only about 20. That one fine-tuning choice cuts concurrent adapter capacity by a factor of about twelve.

Work the KV side too. Llama-3-8B stores 2 x 32 layers x 1024 values x 2 bytes = 128 KB of KV per token. Each extra gigabyte given to adapters is roughly 8,000 fewer tokens of KV cache, which is several fewer concurrent long-context requests. Raise the adapter fraction only when measured churn says the active set is too small, not pre-emptively.

Merging adapters per request

LoRAX can also combine adapters per request. Instead of adapter_id, the request carries merged_adapters with a list of ids, relative weights and a merge_strategy: linear (the default), ties, dare_linear or dare_ties. TIES and DARE also take density, the fraction of adapter weights to keep, and TIES takes majority_sign_method.

{"inputs": "...",
 "parameters": {
   "max_new_tokens": 128,
   "merged_adapters": {
     "ids": ["acme/tone-formal", "acme/domain-insurance"],
     "weights": [0.4, 1.0],
     "merge_strategy": "ties",
     "density": 0.2,
     "majority_sign_method": "total"}}}

Merging is useful for quick experiments, such as combining a style adapter with a domain adapter. Treat a merge as a new model for evaluation purposes: merge methods trade off interference between adapters differently, and a merge that looks fine on one prompt can regress another task badly. If a merge goes to production, train or merge it offline, evaluate it, and serve it under its own ID.

Failure modes

• Base mismatch. An adapter trained on a different revision of the base model loads without error and produces degraded output. Record the base model hash in adapter metadata and reject mismatches in your registry before they reach the server.
• Cold-start tails. Long-tail tenants hit cold fetches on nearly every request. Pre-warm adapters with scheduled traffic or keep them in a local volume instead of fetching from the Hub at request time.
• Churn. More hot adapters than active slots or adapter memory make the scheduler evict and reload in a loop. Watch per-adapter queue time and load counts; fix it by sharding tenants across replicas by adapter, not round-robin.
• KV starvation. An oversized adapter reservation shrinks the KV cache, lowering batch size and throughput for everyone. The symptom is preemption or queueing even when few adapters are in use.
• Noisy neighbour. One tenant's flood of long prompts occupies batch slots. Add per-adapter rate limits at the gateway; the server's fairness is between queues, not a quota system.
• Supply chain. Loading adapters by arbitrary ID from a public hub lets callers choose which weights run on your hardware. Allow-list sources and IDs, and prefer safetensors files.

Trade-offs

What to do next

1. Inventory your fine-tunes: base model revision, rank, target modules and size of each, and drop any adapter whose base does not match the deployment.
2. Compute adapter size and KV bytes per token for your base model as in the worked example, and choose --adapter-memory-fraction from that, not from the default.
3. Deploy one pinned LoRAX image with a local or S3 adapter store and an allow-list of IDs.
4. Load-test with your real traffic distribution across adapters and record TTFT by adapter temperature: resident, host-cached and cold.
5. Add dashboards for per-adapter queue time, load count and preemptions, and alert on churn.
6. Promote any production merge to its own evaluated adapter ID instead of merging per request. Background reading: LoRA fine-tuning, continuous batching and paged KV cache.