← Back to DevBytes

Serverless LLM Inference: Cold Start Optimization Strategies

Serverless LLM Inference: Cold Start Optimization Strategies

Serverless computing has transformed how developers deploy applications, offering automatic scaling, pay-per-use billing, and zero infrastructure management. However, when applied to Large Language Model (LLM) inference, serverless platforms introduce a significant challenge: cold starts. Unlike traditional web applications where a cold start might add a few hundred milliseconds of latency, an LLM cold start can take anywhere from 10 seconds to several minutes. This happens because the platform must provision compute resources, download or load a multi-gigabyte model, initialize the inference engine, and warm up the runtime before serving the first request.

This tutorial explores practical strategies to minimize cold start times for serverless LLM inference, with hands-on code examples you can apply to your own deployments.

What Is a Serverless LLM Cold Start?

A cold start occurs when a serverless function is invoked after being idle or when scaling out to a new instance. For LLM inference, the cold start lifecycle involves several stages:

Each stage contributes to the total cold start time. A 7B parameter model in FP16 format is roughly 14 GB, and loading it into GPU memory alone can take 20-60 seconds depending on storage bandwidth. This makes cold start optimization not just a nice-to-have but a critical requirement for production LLM services.

Why Cold Start Optimization Matters

Cold starts directly impact user experience and business outcomes. When a user sends a chat message and waits 30 seconds before seeing the first token, they are likely to abandon the interaction. For real-time applications such as conversational AI, code assistants, and search augmentation, latency is the primary quality metric.

Beyond user experience, cold starts affect cost. Many teams work around cold starts by keeping functions warm with periodic ping requests, effectively running always-on infrastructure while paying serverless premiums. Others over-provision concurrency, leading to wasted GPU resources. Optimizing cold starts lets you enjoy the true benefits of serverless: scale-to-zero and pay-per-request.

Cold starts also complicate autoscaling. When traffic spikes, new instances must be provisioned, and each one incurs a cold start penalty. If cold starts are long, requests queue up, timeouts occur, and cascading failures can follow. Reducing cold start time from 60 seconds to 5 seconds fundamentally changes how aggressively you can scale.

Strategy 1: Optimize Container Images

The container image is the first thing loaded during a cold start. Bloated images with unnecessary dependencies, large base layers, and unoptimized model storage add seconds or even minutes to startup time. The goal is to produce the smallest possible image that contains only what is needed for inference.

Use Slim Base Images

Start with a minimal base image. Avoid full Ubuntu or Debian images when a slim variant or specialized runtime image will do. For GPU workloads, use the official NVIDIA CUDA runtime images, but choose the runtime (not devel) variants.

# Dockerfile - Optimized LLM inference image
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PIP_NO_CACHE_DIR=1

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 python3-pip && \
    rm -rf /var/lib/apt/lists/*

# Install only inference dependencies
RUN pip install --no-cache-dir \
    vllm==0.4.2 \
    transformers==4.41.0 \
    fastapi==0.111.0 \
    uvicorn==0.30.0

COPY serve.py /app/serve.py
COPY entrypoint.sh /app/entrypoint.sh

WORKDIR /app
ENTRYPOINT ["./entrypoint.sh"]

Multi-Stage Builds to Strip Build Tools

Use multi-stage builds to compile dependencies in one stage and copy only the artifacts to the final image. This removes compilers, headers, and intermediate build files from the production image.

# Build stage
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder

RUN apt-get update && apt-get install -y python3.10 python3-pip git
RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu121
RUN pip install --no-cache-dir vllm==0.4.2

# Runtime stage
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
COPY --from=builder /usr/local/bin /usr/local/bin

COPY serve.py /app/serve.py
WORKDIR /app
CMD ["python3", "serve.py"]

Store Model Weights Efficiently

Bundling model weights inside the container image makes the image large but eliminates remote download time during cold start. This is a tradeoff: larger images take longer to pull but avoid network latency and potential download failures. For serverless platforms with fast image pull (like AWS Lambda with container image support or Google Cloud Run), bundling weights is often the better choice.

# Bundle quantized model weights directly in the image
COPY ./models/llama-3-8b-instruct-awq/ /models/llama-3-8b-instruct-awq/

# In serve.py, load from local path
MODEL_PATH = "/models/llama-3-8b-instruct-awq"

If the image becomes too large, consider storing weights on a high-speed attached volume (like AWS EFS or a persistent disk) that mounts faster than a remote download from object storage.

Strategy 2: Model Quantization

Quantization reduces model size by representing weights with lower precision. A model in FP16 (16-bit float) can be quantized to INT8, INT4, or even lower, dramatically reducing memory footprint and load time. An 8B parameter model in FP16 is about 16 GB; in 4-bit quantization (AWQ or GPTQ), it shrinks to roughly 4-5 GB. This means less data to load from disk to VRAM during cold start.

Using AWQ Quantization with vLLM

from vllm import LLM, SamplingParams

# Load a pre-quantized AWQ model - much faster cold start
# than loading FP16 and quantizing on the fly
llm = LLM(
    model="TheBloke/Llama-2-7B-Chat-AWQ",
    quantization="awq",
    dtype="float16",
    gpu_memory_utilization=0.90,
    max_model_len=4096,
    enforce_eager=False,
)

sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=512,
)

outputs = llm.generate(["Explain cold starts in serverless computing."], sampling_params)
for output in outputs:
    print(output.outputs[0].text)

Quantizing Your Own Model

If you need to quantize a custom fine-tuned model, use the AutoAWQ library. Perform this step offline and store the quantized model, so the serverless function only loads the already-quantized weights.

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "/models/my-finetuned-llama"
quant_path = "/models/my-finetuned-llama-awq"
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model in FP16
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize - do this OFFLINE, not during cold start
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

Quantization typically reduces cold start time by 50-70% for the model loading phase, with minimal quality degradation for 4-bit methods like AWQ and GPTQ on popular model families.

Strategy 3: Provisioned Concurrency and Warm Instances

Most serverless platforms offer a way to keep a minimum number of instances warm. AWS Lambda has Provisioned Concurrency, Google Cloud Run has minimum instances, and Azure Functions has Premium Plan with pre-warmed instances. These features eliminate cold starts for a baseline number of concurrent requests.

AWS Lambda Provisioned Concurrency with Terraform

resource "aws_lambda_function" "llm_inference" {
  function_name = "llm-inference"
  role          = aws_iam_role.lambda_role.arn
  package_type  = "Image"
  image_uri     = "${aws_ecr_repository.llm_repo.repository_url}:latest"
  timeout       = 120
  memory_size   = 10240  # 10 GB memory (Lambda maps this to vCPU)

  ephemeral_storage {
    size = 10240
  }
}

resource "aws_lambda_alias" "prod" {
  name             = "prod"
  function_name    = aws_lambda_function.llm_inference.function_name
  function_version = "$LATEST"
}

resource "aws_lambda_provisioned_concurrency_config" "warm" {
  function_name                     = aws_lambda_function.llm_inference.function_name
  provisioned_concurrent_executions = 3
  qualifier                         = aws_lambda_alias.prod.name
}

Google Cloud Run Minimum Instances

gcloud run deploy llm-inference \
  --image gcr.io/my-project/llm-inference:latest \
  --region us-central1 \
  --min-instances 2 \
  --max-instances 20 \
  --cpu 4 \
  --memory 16Gi \
  --gpu 1 \
  --gpu-type nvidia-l4 \
  --no-cpu-throttling \
  --concurrency 1 \
  --timeout 300

The tradeoff is cost: provisioned instances incur charges even when idle. The key is to right-size the minimum based on your traffic patterns. Use historical data to determine the baseline concurrency that covers the majority of requests, and let the serverless scaler handle bursts above that baseline.

Strategy 4: Lazy Loading and Progressive Initialization

Not all components of your inference server need to be ready before you can start serving requests. Lazy loading defers non-critical initialization until after the first request is handled, or until the component is actually needed. This reduces the time to first token even if total initialization time remains the same.

Deferred Initialization Pattern

import os
import time
from fastapi import FastAPI, Request
from contextlib import asynccontextmanager

app = FastAPI()

# Global references - loaded lazily
llm_engine = None
tokenizer = None
metrics_client = None
feature_store = None

def load_model():
    """Load the core model - required for inference."""
    global llm_engine, tokenizer
    from vllm import LLM, SamplingParams
    from transformers import AutoTokenizer

    model_path = os.environ.get("MODEL_PATH", "/models/llama-3-8b-awq")
    llm_engine = LLM(
        model=model_path,
        quantization="awq",
        gpu_memory_utilization=0.85,
        max_model_len=4096,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_path)

def load_optional_components():
    """Load non-critical components in the background."""
    global metrics_client, feature_store
    try:
        from prometheus_client import start_http_server
        start_http_server(9090)
        metrics_client = True
    except Exception:
        pass

    try:
        # Connect to feature store for RAG
        import redis
        feature_store = redis.Redis(
            host=os.environ.get("REDIS_HOST", "localhost"),
            port=6379,
            decode_responses=True,
        )
        feature_store.ping()
    except Exception:
        feature_store = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load only the model during startup
    load_model()
    print("Model loaded, ready to serve requests")

    # Defer optional components to a background thread
    import threading
    threading.Thread(target=load_optional_components, daemon=True).start()

    yield

app = FastAPI(lifespan=lifespan)

@app.post("/generate")
async def generate(request: Request):
    body = await request.json()
    prompt = body.get("prompt", "")
    max_tokens = body.get("max_tokens", 256)

    from vllm import SamplingParams
    params = SamplingParams(temperature=0.7, max_tokens=max_tokens)
    outputs = llm_engine.generate([prompt], params)

    return {"response": outputs[0].outputs[0].text}

Pre-warming the Inference Engine

Even after the model is loaded, the first inference request is often slower because the inference engine compiles CUDA kernels, allocates KV cache memory, and warms up internal buffers. Send a dummy request during startup to absorb this penalty before real traffic arrives.

def warm_up_engine(llm_engine):
    """Run a dummy inference to compile kernels and warm caches."""
    from vllm import SamplingParams
    warmup_params = SamplingParams(temperature=0.0, max_tokens=1)
    llm_engine.generate(["Warmup"], warmup_params)
    print("Engine warmed up")

# Call this right after load_model() in the lifespan handler
warm_up_engine(llm_engine)

Strategy 5: Model Sharding and Memory Mapping

For larger models (30B+ parameters) that cannot fit on a single GPU, sharding across multiple devices is necessary. The way you load shards affects cold start time. Instead of loading each shard sequentially, use parallel loading and memory-mapped files to overlap I/O with computation.

Parallel Shard Loading

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import concurrent.futures
import os

def load_shard(shard_path, device_map):
    """Load a single model shard onto a specific device."""
    return torch.load(shard_path, map_location=device_map)

def load_model_parallel(model_path, num_gpus=2):
    """Load model shards in parallel across GPUs."""
    shard_dir = os.path.join(model_path, "pytorch_model")
    shard_files = sorted([
        f for f in os.listdir(shard_dir) if f.startswith("pytorch_model-")
    ])

    device_map = {i: f"cuda:{i}" for i in range(num_gpus)}

    # Load all shards concurrently
    shards = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_gpus) as executor:
        futures = {}
        for i, shard_file in enumerate(shard_files):
            gpu_id = i % num_gpus
            future = executor.submit(
                load_shard,
                os.path.join(shard_dir, shard_file),
                f"cuda:{gpu_id}"
            )
            futures[future] = shard_file

        for future in concurrent.futures.as_completed(futures):
            shard_name = futures[future]
            shards[shard_name] = future.result()

    return shards

# For most use cases, use the built-in device_map="auto"
# which already implements smart sharding
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,  # Critical: avoids loading full model on CPU first
)

The low_cpu_mem_usage=True flag is particularly important. Without it, the entire model is loaded into system RAM before being moved to GPU, doubling memory requirements and adding significant time. With it, weights are loaded directly into GPU memory shard by shard.

Memory-Mapped Weight Loading

Using mmap for weight files allows the operating system to lazily page in model weights as they are accessed, rather than reading the entire file upfront. The safetensors format supports this natively.

from safetensors.torch import load_file
import mmap

# safetensors uses mmap by default when loading
# This means the OS loads pages on demand, reducing initial load time
state_dict = load_file("/models/llama-3-8b/model.safetensors")

# For custom loading with explicit mmap
def load_with_mmap(file_path):
    with open(file_path, "rb") as f:
        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
        # Process memory-mapped data
        return mm

Strategy 6: Caching and Request Coalescing

Caching is a powerful cold start mitigation strategy. If multiple users send similar or identical prompts, caching the response eliminates the need for inference entirely. Even partial caching of prompt prefixes (prompt caching) can significantly reduce compute time after a cold start.

Prefix Caching with vLLM

vLLM supports automatic prefix caching, which caches the KV states of common prompt prefixes. This is especially valuable for system prompts that are shared across all requests.

from vllm import LLM, SamplingParams

llm = LLM(
    model="/models/llama-3-8b-awq",
    quantization="awq",
    enable_prefix_caching=True,  # Enable KV cache reuse for shared prefixes
    gpu_memory_utilization=0.85,
    max_model_len=4096,
)

system_prompt = (
    "You are a helpful coding assistant. "
    "Always provide concise, accurate answers with code examples."
)

# First request - computes and caches the system prompt KV
params1 = SamplingParams(temperature=0.7, max_tokens=256)
response1 = llm.generate([f"{system_prompt}\n\nUser: How do I read a file in Python?"], params1)

# Second request - reuses cached system prompt KV, only computes the new tokens
params2 = SamplingParams(temperature=0.7, max_tokens=256)
response2 = llm.generate([f"{system_prompt}\n\nUser: How do I write a file in Python?"], params2)

print("Response 1:", response1[0].outputs[0].text)
print("Response 2:", response2[0].outputs[0].text)

Response-Level Caching with Redis

import hashlib
import json
import redis
import os

redis_client = redis.Redis(
    host=os.environ.get("REDIS_HOST", "localhost"),
    port=6379,
    decode_responses=True,
)

def get_cache_key(prompt, model_id, params):
    """Generate a deterministic cache key from prompt and parameters."""
    cache_input = json.dumps({
        "prompt": prompt,
        "model": model_id,
        "temperature": params.get("temperature", 0.7),
        "max_tokens": params.get("max_tokens", 256),
        "top_p": params.get("top_p", 1.0),
    }, sort_keys=True)
    return f"llm_cache:{hashlib.sha256(cache_input.encode()).hexdigest()}"

def cached_generate(prompt, model_id, params, generate_fn, ttl=3600):
    """Wrapper that checks cache before calling the model."""
    cache_key = get_cache_key(prompt, model_id, params)

    # Check cache first
    cached = redis_client.get(cache_key)
    if cached:
        return json.loads(cached)

    # Cache miss - call the actual model
    response = generate_fn(prompt, params)

    # Store in cache with TTL
    redis_client.setex(cache_key, ttl, json.dumps(response))
    return response

# Usage
def actual_generate(prompt, params):
    from vllm import SamplingParams
    sp = SamplingParams(temperature=params["temperature"], max_tokens=params["max_tokens"])
    outputs = llm_engine.generate([prompt], sp)
    return {"text": outputs[0].outputs[0].text}

result = cached_generate(
    prompt="What is serverless computing?",
    model_id="llama-3-8b-awq",
    params={"temperature": 0.0, "max_tokens": 256},
    generate_fn=actual_generate,
    ttl=7200,  # Cache for 2 hours
)

Strategy 7: Platform-Specific Optimizations

AWS Lambda: SnapStart for Faster Initialization

AWS Lambda SnapStart captures a snapshot of the initialized runtime and restores it instead of running the full initialization code. While originally designed for Java, the concept applies to any runtime where initialization is expensive. For Python-based LLM serving, you can use Lambda's container image support with optimized loading.

# Lambda handler with optimized cold start
import json
import os

# Initialize model at module level - runs once per container
# Use lazy loading to avoid blocking if not needed
_model = None

def get_model():
    global _model
    if _model is None:
        from vllm import LLM
        _model = LLM(
            model=os.environ["MODEL_PATH"],
            quantization="awq",
            gpu_memory_utilization=0.85,
            max_model_len=2048,  # Lower context = faster init
        )
    return _model

def lambda_handler(event, context):
    model = get_model()
    from vllm import SamplingParams

    body = json.loads(event.get("body", "{}"))
    prompt = body.get("prompt", "")
    params = SamplingParams(
        temperature=body.get("temperature", 0.7),
        max_tokens=body.get("max_tokens", 256),
    )

    outputs = model.generate([prompt], params)
    response = outputs[0].outputs[0].text

    return {
        "statusCode": 200,
        "body": json.dumps({"response": response}),
    }

Google Cloud Run: CPU Always Allocated

Cloud Run by default throttles CPU when no request is being processed. For LLM inference, this means the model gets unloaded or the process gets suspended between requests, effectively causing a cold start on every invocation. Enable CPU always allocated to keep the model resident in memory.

gcloud run deploy llm-inference \
  --image gcr.io/my-project/llm-inference \
  --cpu 4 \
  --memory 16Gi \
  --gpu 1 \
  --gpu-type nvidia-l4 \
  --no-cpu-throttling \
  --min-instances 1 \
  --max-instances 10 \
  --concurrency 4

Kubernetes with KEDA: Scale-to-Zero with Fast Warm-Up

For teams running on Kubernetes, KEDA (Kubernetes Event-Driven Autoscaling) provides serverless-like scale-to-zero with more control over the warm-up process. You can use pre-scaled warm pools and fast image pulling with local registries.

apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: llm-inference-scaler
  namespace: llm-serving
spec:
  scaleTargetRef:
    name: llm-inference-deployment
  minReplicaCount: 1
  maxReplicaCount: 10
  pollingInterval: 10
  cooldownPeriod: 60
  triggers:
  - type: prometheus
    metadata:
      serverAddress: http://prometheus:9090
      metricName: llm_request_queue_depth
      threshold: "2"
      query: llm_request_queue_depth

Best Practices

Putting It All Together: A Complete Serverless LLM Server

Here is a complete, production-ready FastAPI server that combines multiple optimization strategies:

import os
import time
import logging
import threading
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global state
llm_engine = None
redis_client = None
startup_time = None

def load_model():
    """Load the quantized model with optimized settings."""
    global llm_engine
    from vllm import LLM

    start = time.time()
    llm_engine = LLM(
        model=os.environ.get("MODEL_PATH", "/models/llama-3-8b-awq"),
        quantization="awq",
        dtype="float16",
        gpu_memory_utilization=float(os.environ.get("GPU_MEM_UTIL", "0.85")),
        max_model_len=int(os.environ.get("MAX_MODEL_LEN", "4096")),
        enable_prefix_caching=True,
        enforce_eager=False,
    )
    logger.info(f"Model loaded in {time.time() - start:.2f}s")

def warm_up():
    """Send a dummy request to compile kernels and warm caches."""
    from vllm import SamplingParams
    start = time.time()
    llm_engine.generate(["Warmup"], SamplingParams(max_tokens=1))
    logger.info(f"Warmup completed in {time.time() - start:.2f}s")

def init_redis():
    """Initialize Redis connection for response caching."""
    global redis_client
    try:
        import redis
        redis_client = redis.Redis(
            host=os.environ.get("REDIS_HOST", "localhost"),
            port=int(os.environ.get("REDIS_PORT", "6379")),
            decode_responses=True,
            socket_timeout=2,
        )
        redis_client.ping()
        logger.info("Redis connected")
    except Exception as e:
        logger.warning(f"Redis unavailable: {e}")
        redis_client = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global startup_time
    startup_time = time.time()

    # Critical path: load model and warm up
    load_model()
    warm_up()

    # Non-critical: cache client loads in background
    threading.Thread(target=init_redis, daemon=True).start()

    logger.info(f"Total startup: {time.time() - startup_time:.2f}s")
    yield

app = FastAPI(lifespan=lifespan)

@app.get("/health")
async def health():
    return {
        "status": "ready" if llm_engine else "loading",
        "startup_time": startup_time,
    }

@app.post("/generate")
async def generate(request: Request):
    body = await request.json()
    prompt = body.get("prompt", "")
    temperature = body.get("temperature", 0.7)
    max_tokens = body.get("max_tokens", 256)
    stream = body.get("stream", False)

    from vllm import SamplingParams
    params = SamplingParams(
        temperature=temperature,
        max_tokens=max_tokens,
    )

    if stream:
        async def stream_response():
            for output in llm_engine.generate([prompt], params, use_tqdm=False):
                for token in output.outputs[0].text:
                    yield f"data: {json.dumps({'token': token})}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(stream_response(), media_type="text/event-stream")

    outputs = llm_engine.generate([prompt], params)
    return {"response": outputs[0].outputs[0].text}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)

Conclusion

Cold start optimization for serverless LLM inference is not a single technique but a combination of strategies applied across the entire initialization pipeline. By optimizing container images, leveraging model quantization, using provisioned concurrency judiciously, implementing lazy loading and warm-up routines, employing efficient model loading with memory mapping and parallel shard loading, and adding caching layers, you can reduce cold start times from minutes to single-digit seconds. The key is to measure each stage of your cold start, identify the dominant bottleneck, and apply the appropriate optimization. As serverless platforms continue to evolve with features like SnapStart, faster image pulling, and GPU-aware scaling, the gap between serverless and always-on LLM serving will continue to close, making true pay-per-request LLM inference a practical reality for production workloads.

— Ad —

Google AdSense will appear here after approval

← Back to all articles