Wrapping MLX Behind a FastAPI Service

Why graduate from the built-in server

mlx-lm's built-in mlx_lm.server (covered in lm.lesson6) is great for development, demos, and any single-user, low-concurrency use. The moment you need real concurrent request handling, request queueing, authentication, custom logging, or any operational concern beyond "a model that responds to OpenAI-shaped requests," you're better off wrapping mlx-lm in a FastAPI service you control.

This isn't a complicated wrapper — it's about 100 lines of FastAPI code that load a model once at startup and expose /generate and /health endpoints. The win is that you own every operational decision and can swap, fork, or extend any of them without fighting the built-in server's defaults.

The pattern

Load the model once at startup — FastAPI's lifespan context manager handles this. The (model, tokenizer) pair lives in app state, reused across all requests.
Warm up — do one throwaway forward pass during startup. The first inference call after a fresh load incurs MLX's JIT compilation cost; warming up means the first real request doesn't pay this tax.
Expose /generate — accept a JSON body with prompt + sampling params; return generated text. Add streaming via Server-Sent Events for token-by-token delivery.
Expose /health — a cheap endpoint your load balancer or process supervisor can hit to verify the service is alive.
Run with uvicorn — single process, single worker for MLX (concurrency comes from async I/O, not from worker pools, because the GPU is the bottleneck and you can't parallelize across workers efficiently).

What you DON'T do

Don't run multiple worker processes for the same model — they'd each load their own copy and fight for GPU memory. Don't try to parallelize generation across requests at the framework level — MLX serializes GPU access anyway, and Python's GIL plus async I/O is sufficient orchestration. Don't add a request queue at the FastAPI level for any normal traffic — uvicorn's connection handling is already a queue.

The minimum FastAPI service

The code block below is the whole service in one file. Save as app.py, run with uvicorn app:app --host 0.0.0.0 --port 8000, hit POST /generate with a JSON body. Production patterns (rate limiting, auth, structured logging) are deliberately absent — add them as you need them, but the bones are here.

Code

Minimum FastAPI service for MLX (app.py)·python

# Save as app.py; run with: uvicorn app:app --host 0.0.0.0 --port 8000

from contextlib import asynccontextmanager
from fastapi import FastAPI
from pydantic import BaseModel
from mlx_lm import load, generate

# Module-level cache for the loaded (model, tokenizer) pair
state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: load the model once
    print("Loading model...")
    state["model"], state["tokenizer"] = load("mlx-community/Llama-3.2-1B-Instruct-4bit")

    # Warm up — pay the JIT cost during startup, not during the first real request
    _ = generate(state["model"], state["tokenizer"], prompt="warmup", max_tokens=1)
    print("Model loaded and warm.")

    yield   # App is alive here

    # Shutdown: nothing special needed; MLX cleans up at process exit
    state.clear()


app = FastAPI(lifespan=lifespan)


class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 100


@app.get("/health")
async def health():
    return {"status": "ok", "model_loaded": "model" in state}


@app.post("/generate")
async def gen(req: GenerateRequest):
    text = generate(
        state["model"], state["tokenizer"],
        prompt=req.prompt, max_tokens=req.max_tokens, verbose=False,
    )
    return {"text": text}

Hit it with curl·bash

# Health check
curl http://localhost:8000/health
# {"status":"ok","model_loaded":true}

# Generate
curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "Capital of France?", "max_tokens": 20}'
# {"text":"Paris is the capital of France..."}

Exercise

Save the FastAPI service from the code block as app.py. Run uvicorn app:app --host 0.0.0.0 --port 8000 in one terminal; hit /health and /generate from another with curl. Time the first /generate call vs the second — they should be nearly identical because of the warmup step. Two sentences on what you observed and what feature you'd add next.