C.W.K.
Stream
Lesson 04 of 06 · published

Multiple Models, Batch Inference, Sharing Memory

~14 min · batching, multi-model, throughput

Level 0Curious
0 XP0/51 lessons0/15 achievements
0/100 XP to next level100 XP to go0% complete

Two patterns that look similar but aren't

Once you're past the single-model-single-request demo, two production patterns come up. They look related but optimize for different things, and using the wrong one wastes resources.

  • Multiple models in one process — your service hosts several different models (e.g. a 7B for chat + a small embedder + a VLM for image understanding). Memory is shared at the process level; you save the per-Python-process overhead but you compete for unified memory between the models.
  • Batched inference for one model — your service receives many concurrent prompts for the same model and processes them as a batch. Throughput goes up; per-request latency goes up too; the trade is for total compute efficiency.

This lesson covers both patterns and when each is the right call.

Multiple models — the memory accounting

Hosting two models in one process means both their weights live in unified memory simultaneously. The per-model napkin math from foundations.lesson4 just adds up. A 7B Q4 (~5 GB) + a small embedding model (~1 GB) + a 2B VLM (~2 GB) = ~8 GB resident, plus KV cache for whichever is currently generating. Accounting matters more here than in the single-model case because you can't afford either model to suddenly need its full inference budget.

The pattern that works: load all models at startup, keep them resident, route requests to the right model based on the request type. The pattern that doesn't work: load on demand and unload on idle — model load latency (often seconds for larger models) dwarfs inference latency for any individual request.

Batched inference — when throughput beats latency

If you have many concurrent prompts for the same model, batching them processes more tokens per GPU dispatch and improves throughput-per-dollar significantly. The trade is per-request latency: each request waits for the batch window to close before it sees first-token output.

mlx-lm doesn't have first-class batched-generation primitives in the way some PyTorch serving stacks do, so the practical pattern is async request collection at the service layer (FastAPI), then single-prompt inference per request via mlx-lm. For workloads dominated by many concurrent short prompts, look at higher-level serving stacks built on MLX (or wait for mlx-lm's batching APIs to mature).

The decision tree

  • Single model, single user — straight FastAPI service from prod.lesson1, no batching needed.
  • Multiple models, one process, low concurrency — load all at startup, route by request type, accept that you're spending unified memory on availability.
  • Single model, very high concurrency — explore batched serving (Ollama or a custom batching wrapper); accept higher per-request latency for higher throughput.
  • Multiple models AND high concurrency — usually a sign you're outgrowing single-Mac serving; evaluate whether a fleet of single-purpose Macs is more economical.

Code

Multiple models in one FastAPI process·python
# Sketch of a multi-model service (extends prod.lesson1's pattern).
from contextlib import asynccontextmanager
from fastapi import FastAPI
from pydantic import BaseModel
from mlx_lm import load, generate

models: dict = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Loading models...")
    models["chat"], models["chat_tok"] = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
    models["small"], models["small_tok"] = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
    # Warm up each
    for name in ["chat", "small"]:
        _ = generate(models[name], models[f"{name}_tok"], prompt="warmup", max_tokens=1)
    print("All models loaded and warm.")
    yield
    models.clear()

app = FastAPI(lifespan=lifespan)

class Req(BaseModel):
    model: str        # "chat" or "small"
    prompt: str
    max_tokens: int = 100

@app.post("/generate")
async def gen(req: Req):
    if req.model not in ("chat", "small"):
        return {"error": "unknown model"}
    text = generate(
        models[req.model], models[f"{req.model}_tok"],
        prompt=req.prompt, max_tokens=req.max_tokens, verbose=False,
    )
    return {"text": text, "model": req.model}
Async batch collection (lightweight pattern)·python
# A thin batching pattern: collect prompts that arrive within a short window,
# process serially via mlx-lm. Real batched generation is more involved;
# this approximation is good enough when you control your traffic shape.
import asyncio
from collections import deque

class BatchCollector:
    def __init__(self, window_ms: int = 25):
        self.window = window_ms / 1000.0
        self.queue: deque = deque()
        self.lock = asyncio.Lock()

    async def submit(self, prompt: str):
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        async with self.lock:
            self.queue.append((prompt, future))
        await asyncio.sleep(self.window)
        # In a full implementation, drain the queue and process as a batch.
        # The minimum sketch: serial inference, returning each result.
        return await future

External links

Exercise

Take the multi-model FastAPI sketch from this lesson and run it locally. Hit it with two different prompts, one routed to chat and one to small. Note the per-request latency for each — the smaller model should be noticeably faster. Then check your memory utilization (Activity Monitor or vm_stat) — both models should be resident. Two sentences on what the multi-model overhead cost you and whether it's worth it for your workflow.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.