predibase/lorax

Multi-LoRA inference server that scales to 1000s of fine-tuned LLMs

Structural Verdict

How Data Flows Through the System

Requests flow from clients through the router to the inference server, where models and adapters are dynamically loaded and applied during generation

1. Request Reception — HTTP/gRPC requests received with adapter specifications (config: adapter_id, adapter_source, merged_adapters)
2. Validation — Tokenization and parameter validation by ValidationWorker (config: max_input_length, max_new_tokens)
3. Adapter Loading — Dynamic loading of LoRA adapters from hub, local, or S3 sources (config: adapter_source, api_token)
4. Batch Formation — Heterogeneous batching of requests across different adapters (config: max_batch_total_tokens, max_concurrent_requests)
5. Inference — Model forward pass with SGMV LoRA application and flash attention (config: do_sample, temperature, top_k +1)
6. Generation — Token generation with stopping criteria and streaming support (config: stop, max_new_tokens, stream)

System Behavior

How the system actually operates at runtime — where data accumulates, what loops, what waits, and what controls what.

Data Pools

Feedback Loops

• Adapter Scheduling (polling, balancing) — Trigger: adapter_cycle_time_s timer. Action: Check adapter usage and prefetch popular adapters. Exit: Server shutdown.
• Dynamic Batching (auto-scale, balancing) — Trigger: Queue depth and waiting time thresholds. Action: Adjust batch size based on current load. Exit: Optimal batch size reached.
• Memory Management (cache-invalidation, balancing) — Trigger: GPU memory pressure. Action: Offload least recently used adapters to CPU. Exit: Memory usage within limits.

Delays & Async Processing

• Adapter Loading (async-processing, ~variable (depends on adapter size and source)) — First request to new adapter has higher latency
• Batch Formation (batch-window, ~waiting_served_ratio * average_serving_time) — Requests wait for optimal batch formation
• Model Prefill (async-processing, ~proportional to input length) — Initial tokens take longer than decode phase
• Adapter Exchange (scheduled-job, ~adapter_cycle_time_s) — Periodic adapter prefetching and offloading

Control Points

• Max Active Adapters (threshold) — Controls: Number of simultaneously loaded adapters. Default: 128
• Waiting Served Ratio (threshold) — Controls: Balance between latency and throughput in batching. Default: 0.3
• Max Batch Total Tokens (threshold) — Controls: Maximum tokens processed in a single batch. Default: 4096
• Adapter Cycle Time (runtime-toggle) — Controls: Frequency of adapter prefetching and scheduling. Default: 2s
• Use SGMV (feature-flag) — Controls: Whether to use SGMV kernels for LoRA operations. Default: true

Technology Stack

Key Components

• Client (class) — Main Python client for making requests to LoRAX server with adapter support clients/python/lorax/client.py
• ShardedClient (class) — gRPC client for communicating with distributed inference shards router/client/src/sharded_client.rs
• Infer (service) — Core inference orchestration handling request batching and adapter scheduling router/src/infer.rs
• AdapterLoader (service) — Manages dynamic loading and caching of LoRA adapters from various sources router/src/loader.rs
• FlashCausalLM (class) — Optimized causal language model implementation with flash attention and LoRA support server/lorax_server/models/flash_causal_lm.py
• LoraLinear (class) — Linear layer with dynamic LoRA adapter application using SGMV kernels server/lorax_server/layers/lora.py
• BatchingQueue (class) — Queue management for heterogeneous batching across different adapters router/src/queue.rs
• Scheduler (service) — Schedules adapter prefetching and request batching for optimal throughput router/src/scheduler.rs
• MedusaModel (class) — Speculative decoding implementation for faster inference with LoRA adapters server/lorax_server/models/medusa.py
• ValidationWorker (class) — Validates and preprocesses incoming requests with tokenization router/src/validation.rs
• Parameters (type-def) — Configuration parameters for generation including adapter settings clients/python/lorax/types.py

Sub-Modules

Configuration

clients/python/lorax/types.py (python-pydantic)

• ids (List[str], unknown)
• weights (List[float], unknown)
• merge_strategy (Optional[str], unknown) — default: None
• density (float, unknown)
• majority_sign_method (Optional[str], unknown) — default: None

clients/python/lorax/types.py (python-pydantic)

• adapter_id (Optional[str], unknown) — default: None
• adapter_source (Optional[str], unknown) — default: None
• merged_adapters (Optional[MergedAdapters], unknown) — default: None
• api_token (Optional[str], unknown) — default: None
• do_sample (bool, unknown) — default: False
• max_new_tokens (Optional[int], unknown) — default: None
• ignore_eos_token (bool, unknown) — default: False
• repetition_penalty (Optional[float], unknown) — default: None
• +14 more parameters

clients/python/lorax/types.py (python-pydantic)

• inputs (str, unknown)
• parameters (Optional[Parameters], unknown) — default: None
• stream (bool, unknown) — default: False

clients/python/lorax/types.py (python-pydantic)

• inputs (List[str], unknown)
• parameters (Optional[Parameters], unknown) — default: None
• stream (bool, unknown) — default: False

Science Pipeline

1. Tokenization — Convert text to token IDs using HuggingFace tokenizer [string → (sequence_length,)] router/src/validation.rs
2. Embedding Lookup — Token IDs to embeddings via embedding layer [(batch_size, sequence_length) → (batch_size, sequence_length, hidden_size)] server/lorax_server/models/flash_causal_lm.py
3. LoRA Application — Apply adapter weights using SGMV kernel: x + (x @ A @ B) * scaling [(batch_size, sequence_length, hidden_size) → (batch_size, sequence_length, hidden_size)] server/lorax_server/layers/lora.py
4. Attention — Multi-head attention with flash attention optimization [(batch_size, sequence_length, hidden_size) → (batch_size, sequence_length, hidden_size)] server/lorax_server/models/flash_causal_lm.py
5. Token Generation — Logits to next token via sampling or greedy selection [(batch_size, vocab_size) → (batch_size, 1)] server/lorax_server/models/flash_causal_lm.py

Assumptions & Constraints

• [warning] Assumes input tensor is (batch_size, sequence_length, hidden_size) but no explicit shape validation (shape)
• [critical] Assumes all tensors are on CUDA device without device checking (device)
• [info] LoRA rank (r) assumed to be positive integer but no validation enforced (value-range)
• [critical] SGMV kernels expect specific tensor layouts and dtypes (float16/bfloat16) without runtime checks (format)

Related Ml Training Repositories

Frequently Asked Questions

Analyzed on March 31, 2026 by CodeSea. Written by Karolina Sarna.