Inference optimization for DS interviews

Why interviewers care about inference

Training a model is a research exercise. Serving it is the actual product. A Data Scientist who can hand a notebook to MLEs and call it a day will lose to one who can defend p99 latency, GPU utilization, and dollar-per-million-tokens in the same sentence. That gap is exactly what interviewers at Meta, Stripe, Anthropic, and DoorDash probe in the systems round.

The question almost never sounds like "explain ONNX". It sounds like "your recommender returns in 800ms p95, the SLA is 200ms, what do you change?" or "we serve 50k QPS on 8 A100s, costs are blowing up, where do you cut?". To answer either, you need a mental map of the optimization stack — graph compilation, batching policy, numerical precision, and hardware fit — and the order in which they pay off.

This post walks the stack from highest leverage to lowest, with the numbers a senior engineer would actually quote in a loop. Memorize the ranges, not the slides.

ONNX as the portability layer

Open Neural Network Exchange is a standard graph format that lets you train in PyTorch and serve almost anywhere. Export is one line, and the resulting artifact runs under ONNX Runtime on CPU, CUDA, ROCm, CoreML, or DirectML without dragging your training framework into production.

import torch

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    do_constant_folding=True,
)

The interview-grade reason to use ONNX is not "cross-framework cool factor" — it is that ONNX Runtime gives you graph-level optimizations for free: constant folding, operator fusion, layout transforms, and a clean handoff into vendor-specific execution providers like TensorRT or OpenVINO. You write the export once and pick the backend at deploy time.

When you say "we moved to ONNX" in an interview, expect the follow-up: what broke? Answer honestly. Custom ops, dynamic control flow, and aten ops that have not been added to the current opset are the usual suspects. The fix is either a newer opset, a symbolic rewrite, or a small custom op registration — not a framework rewrite.

TensorRT, vLLM, and the runtime zoo

Once your graph is in ONNX, the runtime you pick determines the speedup. Here is the cheat sheet interviewers expect you to internalize.

The interview trap here is treating these as interchangeable. They are not. TensorRT wins on raw latency for vision models with stable input shapes. vLLM wins on throughput for LLMs because its PagedAttention and continuous batching keep the GPU saturated even when prompts have wildly different lengths. Triton is not a runtime — it is the layer above that gives you HTTP/gRPC, metrics, and model versioning around whatever runtime you actually use.

If you are asked to "serve a 70B model on 4 H100s," the right answer mentions vLLM or TensorRT-LLM by name and explains tensor parallelism — not "we use Flask".

Batching: static, dynamic, continuous

Batching is the single highest-leverage knob in inference serving, and it is also the question candidates flub most. The reason GPUs are expensive is that a single forward pass on a batch of 1 leaves 98% of the silicon idle. Group requests and throughput goes up nearly linearly until you saturate memory bandwidth.

Static batching is the textbook version: wait for batch_size requests, fire one forward pass, return all responses. It is simple and useless in production because the last request in a slow batch waits for the first one. Latency variance explodes.

Dynamic batching, as implemented in Triton, fixes this with two knobs: max_batch_size and max_queue_delay_microseconds. The server flushes whichever fires first. Tune max_queue_delay to about 10–30% of your p99 SLA — usually 5–20ms for a 100ms SLA. Anything more and you pay tail latency for nothing.

Continuous batching (also called in-flight batching, popularized by vLLM and TensorRT-LLM) is the LLM-specific variant. Generation is autoregressive, so different requests finish at different token counts. Static batching would force every request to wait for the slowest one. Continuous batching evicts finished sequences from the batch and slots new ones in on every decode step. The result is 5–20x throughput improvement over naive batching at the same latency budget.

Quantization: FP32 to INT4

Lowering numerical precision is the second-biggest lever, and the most commonly misunderstood. Models train in FP32 because gradients need range. Inference does not need that range — you can usually drop to FP16 or BF16 with no measurable accuracy loss and get a 2x memory and speed win on modern GPUs. That is the freebie.

The interesting trade-offs start at INT8 and below. Two methods matter:

Post-training quantization (PTQ) runs a calibration pass on a small representative dataset, computes per-tensor or per-channel scales, and writes out an INT8 model. It takes minutes and usually costs 0.3–1.5% accuracy on vision and tabular models. For most production systems this is the right default.

Quantization-aware training (QAT) simulates quantization during fine-tuning so the model learns weights that round well. It costs a training run but recovers most of the accuracy gap. Use QAT when PTQ loses more than 1% on your eval set and you cannot afford it.

For LLMs specifically, INT4 with GPTQ or AWQ is the workhorse — a 70B model fits in ~35GB of VRAM instead of 140GB, which is the difference between one H100 and four. Mention this number in interviews; it lands.

The mistake to avoid is quantizing the wrong layers. Embedding tables and the last classification head are usually FP16 even in an INT8 model — they are tiny and very sensitive. A good quantization toolkit (TensorRT, bitsandbytes, AutoGPTQ) handles this; a naive one does not.

Edge inference

Edge means the model runs where the user is — phone, browser, car, smart camera — not in a datacenter. The constraints flip: instead of optimizing for cost per million tokens, you optimize for battery, thermal headroom, and binary size.

The tooling stack is different too. TensorFlow Lite and its successor LiteRT dominate Android. CoreML is the only first-class story on iOS — it compiles to the Apple Neural Engine, which beats the GPU on energy efficiency by 3–5x for transformer workloads. ONNX Runtime Mobile is the cross-platform fallback. For LLMs on consumer hardware, llama.cpp with GGUF quantization is the reference implementation — a 7B model in INT4 runs at 20+ tokens/sec on an M-series Mac with no GPU code at all.

The tactics map to one principle: shrink the model, then squeeze it. Start with a distilled or natively smaller architecture (MobileNet, DistilBERT, Phi-3 mini), apply aggressive quantization (INT4 is the new normal on edge), prune structured sparsity if your hardware supports it, and compile to the exact device with a vendor toolchain. Skipping the distillation step and trying to quantize a 70B model down to a phone is the most common rookie answer in interviews — it does not work, and saying so out loud is worth points.

Common pitfalls

The first pitfall is optimizing the wrong thing. Candidates jump to TensorRT or INT8 before they have profiled. Real bottlenecks are often in tokenization, image decode, or network hops — the model is 30ms and the request is 800ms. The fix is to run a real profiler (Nsight Systems, PyTorch profiler, py-spy) before touching the model. In interviews, lead with "first I'd profile the end-to-end path" and you will outscore 80% of candidates instantly.

A second trap is ignoring batch size when reporting benchmarks. "Our model runs at 5ms" is meaningless without batch size. A model can do 5ms at batch=1 and 8ms at batch=64 — eight milliseconds for sixty-four answers. Always quote throughput and latency together, and tie them to the SLA. The cleanest format is "p50 / p99 latency at the batch size that hits the throughput target."

The third pitfall is mixing up training-time and inference-time accuracy. Quantization, pruning, and ONNX export all introduce tiny graph differences. A model that scored 0.94 AUC in PyTorch can score 0.937 in the deployed ONNX-INT8 build. That is fine if you measure it. Set up an offline eval pipeline that runs the exact serving artifact against a held-out set, and treat any drop greater than your noise floor as a regression. Skipping this step is how you ship a quietly worse model and only notice on Monday's KPI review.

A fourth subtle issue is KV-cache memory math for LLM serving. New candidates underestimate the cost. A 13B model serving 256 concurrent users at 2048 tokens of context can spend more VRAM on KV-cache than on weights. PagedAttention and KV-cache quantization (FP8 or INT8 cache) exist for this reason. If you cannot ballpark cache_bytes = 2 * num_layers * num_heads * head_dim * seq_len * batch * dtype_bytes in your head, you will get caught.

The fifth pitfall is deploying a runtime you have not measured under realistic load. vLLM is fantastic on benchmarks and merciless on misconfigured systems. TensorRT compile times can hit 20+ minutes for large models, which kills your CI if you rebuild on every commit. Always run a load test with the production traffic shape — long prompts, short prompts, ragged batches, occasional cancellations — before declaring victory.

Related reading

• ML latency optimization for DS interviews
• MLOps for DS interviews
• MLOps monitoring for DS interviews
• Canary and shadow deployment for DS interviews
• Hallucinations and LLM evals for DS interviews

If you want to drill DS systems questions like this every day, NAILDD is launching with 1,500+ interview problems across exactly this pattern.

FAQ

Is this official guidance from NVIDIA or vLLM?

No. The numbers and patterns above are synthesized from public docs (ONNX, TensorRT, vLLM, llama.cpp) and standard benchmark suites like MLPerf Inference. Treat the speedup ranges as rough order-of-magnitude anchors — real numbers depend on model, GPU generation, and batch shape.

When should I pick TensorRT over vLLM?

TensorRT is the right answer for vision, recsys, and any model with stable input shapes on NVIDIA GPUs — lowest absolute latency. vLLM wins for LLM serving with variable prompt lengths and high concurrency because PagedAttention and continuous batching solve KV-cache fragmentation. TensorRT-LLM blurs the line and is now competitive for production LLM workloads on H100/Blackwell.

How much accuracy do I lose with INT8 quantization?

For most vision and tabular models, 0.3–1.5% on the target metric after a good calibration pass. For transformer LMs, INT8 weights with FP16 activations is usually within 1% of FP16 perplexity. INT4 with GPTQ or AWQ loses 1–3% perplexity, fine for chat but worth benchmarking for math or code generation.

What is continuous batching, in one paragraph?

In autoregressive generation, different requests finish at different token counts. Static batching forces every request to wait for the slowest one. Continuous batching reorganizes the batch on every decode step: finished sequences exit, new ones slot in, the GPU never idles. The result is 5–20x throughput at the same per-token latency.

Do I need all of this for a Data Scientist role, or is it MLE territory?

For pure modeling DS at smaller companies, knowing ONNX, FP16, and batching exist is enough. For senior DS at scaled product companies (Meta, Stripe, Anthropic, Uber, DoorDash) and anything with "ML Engineer" or "Applied Scientist" in the title, defend the full stack end-to-end — runtime, batching, quantization, KV-cache math, edge constraints. The systems-aware DS gets the offer.