Serving Fine-Tuned Models: Quantization and Inference Optimization | FromZero2AI Blog

Deploying Large Language Models (LLMs) into production environments presents a unique set of engineering challenges. Unlike traditional machine learning models, LLM inference is highly resource-intensive and primarily memory-bandwidth bound. During the autoregressive generation phase, the model must read all its weights from high-bandwidth GPU memory (VRAM) to the local processor cache (SRAM) to compute just a single token.

To make hosting economically viable and reduce latency, modern inference engines employ sophisticated optimizations. This article explores the mechanics of KV Caching, Continuous Batching, and popular quantization schemes like AWQ, GPTQ, and GGUF.

The Key-Value (KV) Cache and Memory Footprint

During autoregressive generation, a transformer model predicts the next token by looking at all preceding tokens in the prompt and the tokens it has generated so far. This requires calculating the Self-Attention matrix:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

For a sequence of length N, the Key (K) and Value (V) matrices of previous tokens do not change. Computing them repeatedly for every new token represents O(N^2) redundant computations.

The KV Cache solves this by storing the K and V matrices of all processed tokens in memory. On subsequent tokens, the model only computes the Q, K, and V vectors for the single new token, appending its K and V to the cache.

Memory Overhead of the KV Cache

While KV caching reduces computational load, it introduces a significant memory footprint. The memory occupied by the KV cache per model (in bytes) is calculated as:

\text{Memory}_{\text{KVCache}} = 2 \times 2 \times N_{\text{layers}} \times N_{\text{heads}} \times d_{\text{head}} \times L_{\text{seq}} \times B_{\text{batch}}

The first factor of 2 accounts for storing both K and V matrices.
The second factor of 2 represents the byte size of 16-bit precision (FP16 or BF16).
N_{\text{layers}} is the number of transformer layers.
N_{\text{heads}} is the number of key-value heads (using Multi-Query or Grouped-Query Attention reduces this).
d_{\text{head}} is the dimension of each head.
L_{\text{seq}} is the sequence length.
B_{\text{batch}} is the batch size.

For a 70B model with 80 layers, 8 KV heads, and a head dimension of 128, a batch size of 32 running a context of 4096 tokens requires approximately 21.5 GB of VRAM just for the KV cache.

Advanced Inference Frameworks

1. PagedAttention

Traditional KV caches allocate contiguous memory blocks for each request. This leads to heavy memory fragmentation (internal and external) because the maximum sequence length must be pre-allocated.

Developed by vLLM, PagedAttention partitions the KV cache into small, fixed-size physical blocks (similar to virtual memory paging in operating systems). The keys and values for a sequence are mapped to non-contiguous blocks, eliminating fragmentation and enabling up to 96% memory reclamation.

2. Continuous Batching

Traditional static batching waits for a batch of requests to arrive, runs them together, and returns the results when the longest generation in the batch completes. This introduces idle time, as some sequences finish early but cannot be released.

Continuous Batching (or iteration-level scheduling) intercepts the execution loop at the token level. As soon as a request in the batch generates its End-Of-Sequence (EOS) token, it is evicted, and a new request is immediately injected into the active execution batch without waiting for other requests to finish.

Quantization Paradigms

Quantization compresses model weights from 16-bit floats to lower-bit representations (e.g., 8-bit or 4-bit) to reduce VRAM requirements and increase memory bandwidth efficiency.

Serving Fine-Tuned Models: Optimized model serving diagram showing KV Caching, PagedAttention blocks, and continuous batching Serving Fine-Tuned Models: Optimized model serving diagram showing KV Caching, PagedAttention blocks, and continuous batching

GPTQ (Generalized Post-Training Quantization)

GPTQ is an accurate one-shot weight quantization method based on approximate second-order information. It processes layer weights by resolving a least-squares problem, correcting the remaining weights to compensate for the quantization error introduced by quantizing each column. GPTQ is ideal for high-throughput GPU serving.

AWQ (Activation-aware Weight Quantization)

AWQ builds on the observation that not all weights in an LLM are created equal. By examining activation distributions during calibration, AWQ identifies the top 1% of "salient" channels that contribute most to accuracy. It protects these salient channels by scaling their weights before quantizing the remaining 99% to 4-bit precision, minimizing accuracy degradation without requiring runtime mixed-precision execution.

GGUF (GGML Unified Format)

Introduced by the llama.cpp project, GGUF is a binary file format designed for fast loading and running of LLMs on consumer-grade hardware. GGUF supports split GPU-CPU execution, allowing layers of the model to be offloaded to CPU system RAM if the GPU VRAM is full.

Serving Quantized Models with vLLM

Below is a Python script demonstrating how to programmatically set up an optimized inference engine using the vLLM library to serve an AWQ-quantized model with continuous batching and PagedAttention.

from vllm import LLM, SamplingParams

# 1. Initialize vLLM Engine
# We load an AWQ quantized model and specify the quantization scheme.
# gpu_memory_utilization controls the fraction of GPU memory allocated for the KV cache.
llm = LLM(
    model="TheBloke/Llama-2-7B-Chat-AWQ",
    quantization="awq",
    tensor_parallel_size=1,  # Number of GPUs to use
    max_model_len=2048,
    gpu_memory_utilization=0.85
)

# 2. Configure Inference Parameters
# temperature controls randomness; max_tokens limits generation length
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=256,
    stop=["[INST]", "User:", "\n\n"]
)

# 3. Define Batch Prompts
prompts = [
    "Explain the concept of quantum computing in simple terms.",
    "Write a short python function that calculates the Fibonacci sequence.",
    "What are the main differences between SQL and NoSQL databases?"
]

# 4. Execute Inference
# vLLM automatically applies PagedAttention and continuous batching under the hood
outputs = llm.generate(prompts, sampling_params)

# 5. Process and Display Outputs
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"\n--- Prompt: {prompt} ---")
    print(f"Response: {generated_text}")

Performance Matrix

Understanding these quantization types and deploying them with frameworks like vLLM enables companies to reduce operational costs by up to 70% while improving user-facing latency.