PaperlayerView exact revision on arXiv

Source-linked AI summary

Punica: Multi-Tenant LoRA Serving

Lequn Chen, Zihao Ye, Yongji Wu, Danyang Zhuo, Luis Ceze, Arvind Krishnamurthy

arXiv:2310.18547v1cs.DCcs.LG

TL;DR

Serving many specialized LoRA models can waste GPUs when each is treated independently, despite models sharing a pretrained base. Punica batches work across different LoRA models and consolidates requests, achieving 12x higher throughput than state-of-the-art LLM serving systems with 2ms added latency per token.

  • Problem

    Serving many tenant-specific LoRA models independently can overlook their shared pretrained weights and require substantial GPU resources.

    Source: §1 · p. 1

  • Method

    Punica combines a CUDA kernel for batching GPU operations across different LoRA models with a scheduler that consolidates multi-tenant workloads.

    Source: §1 · pp. 1–2

  • Results

    12x higher throughput than state-of-the-art LLM serving systems, with only 2ms latency added per token.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A GPU can serve multiple different LoRA models while storing only one copy of the underlying pretrained model, improving memory and computation efficiency.

    Source: §9 · p. 10

Abstract

from arXiv · show

Low-rank adaptation (LoRA) has become an important and popular method to adapt pre-trained models to specific domains. We present Punica, a system to serve multiple LoRA models in a shared GPU cluster. Punica contains a new CUDA kernel design that allows batching of GPU operations for different LoRA models. This allows a GPU to hold only a single copy of the underlying pre-trained model when serving multiple, different LoRA models, significantly enhancing GPU efficiency in terms of both memory and computation. Our scheduler consolidates multi-tenant LoRA serving workloads in a shared GPU cluster. With a fixed-sized GPU cluster, our evaluations show that Punica achieves 12x higher throughput in serving multiple LoRA models compared to state-of-the-art LLM serving systems while only adding 2ms latency per token. Punica is open source at https://github.com/punica-ai/punica .

1 INTRODUCTION

Punica targets the GPU waste of serving each LoRA model independently by batching requests across different LoRA models and consolidating workloads. Its CUDA kernel and scheduler support multi-tenant serving, with 12x higher throughput than state-of-the-art LLM serving systems and 2ms additional latency per token.

1 INTRODUCTION

  • Serving each LoRA model as an independent model can require k × n GPUs for n models needing k GPUs each, overlooking their shared pretrained weights.LoRA retains pretrained weights while adding trainable rank-decomposition matrices, allowing tenants to train specialized models at low cost.Source: §1 · p. 1
  • SGMV batches GPU operations for concurrent, different LoRA models, letting each GPU store only one copy of the pretrained model.The design improves GPU efficiency in both memory and computation.Source: §1 · pp. 1–2
  • SGMV batches different LoRA models with negligible performance differences from batching the same models, while on-demand LoRA loading adds only millisecond-level latency.This lets Punica consolidate requests without being constrained by which LoRA models are already running on a GPU.Source: §1 · p. 1
  • Punica routes new requests to active GPUs until they are fully utilized, then allocates more GPUs; it periodically migrates existing requests for consolidation.The stated design guidelines prioritize consolidating workloads and enabling batching across different LoRA models.Source: §1 · pp. 1–2
  • 12x higher throughput than state-of-the-art LLM serving systems, with only 2ms latency added per token.Evaluated on LoRA models adapted from Llama2 7B, 13B, and 70B on NVIDIA A100 GPU clusters.Source: §1 · p. 2

2 BACKGROUND

Transformer text generation alternates prompt prefill with iterative decoding, whose single-vector inputs leave GPU utilization low. LoRA reduces fine-tuning costs, while serving multiple LoRA models raises questions about shared weights and efficient batching.

2.1 Transformer and Text Generation

  • Decode iteratively consumes a token and KvCache, emits another token, and appends to the cache; its single-vector input results in low GPU utilization.Prefill accepts a prompt and generates a subsequent token and KvCache.Source: §2.1 · p. 2
  • Batching decode requests can improve GPU utilization: increasing batch size from 1 to 32 raises decode latency from 11ms to 13ms for short sequences and 17ms to 34ms for longer sequences.Decode predominately determines serving latency for long-output responses.Source: §2.1 · p. 2

2.2 Low-Rank Adaptation (LoRA)

  • LoRA represents the fine-tuning weight difference as two small dense matrices, reducing trained parameters; each fine-tuned model adds 0.1% to 1% of model weight.LoRA rank is usually much smaller than the original dimension, such as 16 instead of 4096.Source: §2.2 · p. 2
  • Serving each LoRA model independently overlooks shared pre-trained weights, while efficient batching of LoRA add-ons remains an open question.The section frames efficient multi-tenant LoRA serving on a shared GPU cluster as the central question.Source: §2.2 · pp. 2–3

3 PUNICA OVERVIEW

Punica serves multiple LoRA models on shared GPUs by retaining the pre-trained backbone and batching adapter computation. Its SGMV kernel supports this computation, while the scheduler consolidates requests across GPUs.

3 PUNICA OVERVIEW

  • The serving architecture routes user requests through frontends and a scheduler to GPU runners, which stream generated tokens back through the scheduler and frontends.Each request identifies a LoRA model and includes a prompt.Source: §3 · p. 3
  • Punica keeps the pre-trained model loaded on each GPU and swaps in only needed LoRA components, allowing new models to start by loading matrices A and B.A large fraction of GPU memory is reserved for KvCache.Source: §3 · p. 3
  • SGMV batches LoRA add-on computation across requests, grouping requests for the same LoRA model to increase operational intensity and use Tensor Cores.The backbone uses existing matrix multiplication; SGMV computes the LoRA additions.Source: §3 · p. 3
  • Punica schedules requests to active GPUs and periodically migrates older requests to consolidate workloads and free GPU resources.SGMV allows batch-size increases across different LoRA models to improve GPU utilization.Source: §3 · p. 3

4 SEGMENTED GATHER MATRIX-VECTOR MULTIPLICATION

SGMV batches inputs grouped by LoRA model for dense projections. Its expand and shrink cases use different GPU schedules, while the backbone computation remains batched through regular GEMM.

4 SEGMENTED GATHER MATRIX-VECTOR MULTIPLICATION

  • SGMV partitions batch inputs into consecutive groups for each LoRA model before computing the dense projection output.The sequence index s_i marks the last element for model i, with s_0 = 0 and s_n equal to batch size.Source: §4 · p. 3
  • The backbone uses regular batched GEMM, while the LoRA add-on is computed in two kernel launches: first xA, then vB.The intermediate v is initialized to zero before the two launches.Source: §4 · pp. 3–4
  • SGMV-shrink maps high-dimensional input features to low-rank output, while SGMV-expand maps low-rank input features to high-dimensional output.These are the two operators in a LoRA module.Source: §4 · p. 4
  • The expand schedule splits A along the output-feature dimension and dispatches the resulting computations to different GPU threadblocks.The shrink schedule uses Split-K to increase parallelism, then synchronizes and reduces partial outputs across threadblocks.Source: §4 · p. 4
  • When each request has a distinct LoRA index, the computation becomes IO-bound matrix-vector multiplication, so the schedule maximizes memory-bandwidth utilization without Tensor Cores.The passage attributes the Tensor Core choice to the operator's low operational intensity.Source: §4 · p. 4

5 PUNICA IN DETAIL

Punica combines request batching, scheduling, on-demand LoRA loading, migration, and a separable KvCache layout to serve multi-tenant workloads across GPUs. Its design aims to consolidate requests while managing GPU batch capacity and memory.

5 PUNICA IN DETAIL

  • Punica combines prefill and decode requests in one model invocation, limiting each prefill batch to one while batching dense projections and LoRA operations.Requests that finish are removed from the working set; GPUs continuously run prefill and decode steps.Source: §5 · p. 4

5.1 Scheduling new requests

  • The scheduler assigns new requests to the largest eligible working set, subject to batch-size and KvCache-memory limits, and queues requests when no GPU qualifies.The maximum batch size balances cluster throughput and per-token latency; GPU selection also supports workload consolidation and cluster scaling.Source: §5.1 · pp. 4–5

5.2 On-demand model loading

  • LoRA weights load asynchronously while a GPU runs other batch inputs, allowing a newly added request to join after the model invocation completes.Loading the full LoRA model takes around 2ms on PCIe Gen4 x16; decode steps take around 30ms.Source: §5.2 · p. 5

5.3 Request migration

  • When KvCache memory runs low, Punica evicts the newest request and adds it to another GPU, where prefill recomputes its cache from the prompt and generated tokens.The eviction order preserves FCFS semantics; Punica opts for recomputation rather than moving KvCache.Source: §5.3 · p. 5, Figure 5 · p. 5

5.4 Memory layout for KvCache

  • Inseparable KvCache can force shorter requests to run extra decode steps; Punica puts the batching dimension outermost in a paged layout to enable continuous batching.The HuggingFace layout concatenates the cache at each decode step, reading and rewriting it, and keeps requests batched together until all reach their stopping conditions.Source: §5.4 · pp. 5–6, Figure 6 · p. 6

6 IMPLEMENTATION

Punica implements its GPU library as a PyTorch extension and uses supporting libraries and fused operations for model computation. It also constructs batch metadata and SGMV segment indices once per model invocation, alongside Rust-based cluster components.

Python Library

  • Punica exposes CUDA kernels through a PyTorch extension and uses FlashInfer for self-attention and paged KvCache; fused LayerNorm reduces latency from 110µs to 4µs.The Llama implementation is adapted from HuggingFace Transformers.Source: §6 · p. 6
  • Punica groups requests sharing a LoRA model consecutively and builds BatchLen and SGMV segment indices once per model invocation to avoid repeated computation across layers.Prefill requests precede decode requests; the metadata distinguishes prefill start indices and the number of decode requests.Source: §6 · p. 6

Other system components

  • The scheduler, frontend, and runner are written in Rust; each GPU’s runner communicates with a Python subprocess through pipes.Unary RPC and streaming text chunks use web sockets, and I/O is handled asynchronously.Source: §6 · p. 6

7 EVALUATION

Punica is evaluated through kernel, transformer-layer, text-generation, and cluster-deployment tests. Across these settings, it batches requests for different LoRA models and maintains throughput across workload distributions.

7.1 Microbenchmarks

  • In the roofline analysis, SGMV is memory-bandwidth-bound for Identical workloads, while Distinct workloads keep arithmetic intensity constant as batch size increases.Uniform and Skewed workloads combine increasing parallelism and arithmetic intensity; the roofline plot compares arithmetic intensity with achieved throughput.Source: §7.1 · p. 7, Figure 7 · p. 7
  • SGMV outperforms baseline LoRA operators across workloads, with latency rising from 37µs to 116µs in Distinct and only 37µs to 46µs in Uniform and Skewed.In Identical, SGMV latency stays nearly constant from 37µs to 40µs; Gather-BMM incurs more memory I/O than SGMV.Source: §7.1 · pp. 7–8, Figure 8 · p. 8
  • At batch size 64 in Distinct workloads, LoRA ranks 8, 16, 32, and 64 reach 72µs, 75µs, 89µs, and 118µs, respectively.With weight sharing in Uniform, Skewed, and Identical workloads, latency stays around 42µs to 45µs across batch sizes 1 to 64.Source: §7.1 · p. 8, Figure 9 · p. 8
  • At sequence length 512, transformer-layer latency increases only 72% as batch size grows from 1 to 32; longer sequences reduce the batching effect because self-attention takes longer.Layer latency is roughly the same across workloads because the LoRA add-on is small relative to backbone dense projection and self-attention, allowing scheduling to focus on overall throughput.Source: §7.1 · p. 8, Figure 10 · p. 9

7.2 Text generation

  • Punica delivers 1044 tok/s on 7B and 693 tok/s on 13B models, with throughput remaining high across workloads.Most baselines deteriorate when requests target multiple LoRA models; Punica can batch different LoRA models together, maintaining batch size 32.Source: §7.2 · pp. 8–9, Figure 11 · p. 9
  • 441–446 tok/s for Punica versus around 21–25 tok/s for vLLM when serving a 70B model with multiple LoRA models.For the Identical case, Punica and vLLM achieve the same performance; Punica's throughput remains consistent across LoRA popularity distributions.Source: §7.2 · p. 9, Figure 12 · p. 9

7.3 Cluster deployment

  • During a one-hour 16-GPU deployment, GPUs usually reach maximum batch size when active, while some requests migrate when a GPU runs out of KvCache space.Idle GPUs may remain idle and can be released to the cloud provider if necessary.Source: §7.3 · pp. 9–10, Figure 13 · p. 9

8 RELATED WORK

Related work addresses LLM inference optimization, multi-model serving, and memory reduction. Punica combines existing LLM optimizations with LoRA-specific batching and shared pre-trained model weights.

LLM inference optimization

Multi-model inference serving

  • Unlike CNN serving systems, LLM serving must persist KvCache; PetS shares pre-trained model memory across adapters but does not run different models concurrently.Punica runs batches back-to-back due to KvCache affinity, while Nexus has limited support for LLMs and fine-grained sharing patterns.Source: §8 · p. 10

Model and KvCache quantization/compression

  • Model quantization saves memory headroom for KvCache, while KvCache quantization and compression reduce memory I/O and can reduce inference latency.The passage notes that quantization reduces self-attention latency, making Punica's LoRA kernel efficiency more important.Source: §8 · p. 10

9 CONCLUSION

Punica serves multiple LoRA models in a shared GPU cluster using a CUDA kernel that batches operations across models and a scheduler that consolidates workloads. Evaluations report higher throughput with a small latency increase.

9 CONCLUSION

  • With a fixed-sized GPU cluster, Punica achieves 12x higher LoRA serving throughput than state-of-the-art LLM serving systems while adding 2ms latency per token.Each GPU needs only one copy of the underlying pre-trained model to serve multiple LoRA models.Source: §9 · p. 10
Loading 2310.18547v1…