AI Lab System Design Interview

The last chapter practiced the small stateful pieces: queues, rate limiters, cancellation flags, and explicit state machines. A system-design round asks you to place those same pieces on a board and defend the boundaries.

Don't start by drawing every box you know. Name who uses the gateway, size the constraint that actually hurts (often GPU memory or a permission SLO), then add queues, caches, model routing, eval gates, or human review only where that constraint forces them.

We'll keep one running product: an internal model gateway that teams use to call multiple LLM providers. Later follow-ups bolt on permission-aware retrieval and a long-running coding agent, the way an interviewer stretches a first design.

The design round script

Open with a constraint, not a component list:

I'll keep the first design simple, size the constraints early, then add queues, caches, sharding, model routing, or eval gates only when the requirement forces them.

Then follow this order. Each step earns the next: you can't size GPUs before you know the latency SLO, and you shouldn't draw a vector index before you know whether the product retrieves.

1. Goal: who uses it and what success means.
2. Requirements: functional and non-functional.
3. Scale: QPS, tokens, tenants, documents, latency, retention.
4. API: external contract and important internal interfaces.
5. Data model: entities, indexes, isolation, retention.
6. Architecture: simplest request path first.
7. Reliability: retries, idempotency, backpressure, overload, failover.
8. Safety/security: permissions, audit, abuse controls, rollback.
9. Observability: metrics, logs, traces, support views.
10. Rollout: beta gates, canaries, eval gates, kill switches.

Why start with the goal instead of Kafka, GPUs, or vector databases?

Answer

Because the right infrastructure depends on the product goal and constraints. A support lookup, a long-running coding task, and a model-serving scheduler have different latency, reliability, privacy, and eval requirements.

Design pattern taxonomy

Most frontier AI/backend prompts are combinations of these patterns. Identify the dominant pressure before drawing boxes.

Use this table to avoid architecture soup. A model gateway prompt doesn't need a vector database unless the product asks for retrieval. A retrieval prompt doesn't need autonomous agent planning unless the user asks the system to take actions. The running product here is the first four rows: gateway, scheduler, retrieval, and agent.

Follow-up response bank

Interviewers often stress the first design with a new constraint. Answer by naming the boundary you'll change.

Scale math checklist

Every design should include one small calculation. It doesn't need perfect precision; it needs to expose the bottleneck.

Use this phrasing:

I'll size the constraint that most affects the design. If that assumption changes, the architecture boundary I would revisit is X.

Serving math you should have ready

For any serving prompt, two numbers decide the design before you draw a single box: how much GPU memory the weights take, and how fast the KV cache grows per request. You don't need exact figures. You need to show that you know where the memory goes.

The weight footprint is the parameter count times the bytes per parameter:

$M_\text{weights} = P \times b$

Here $P$ is the number of parameters and $b$ is bytes per parameter: 2 for FP16 or BF16, 1 for FP8. An 8B-class model in BF16 is about $8\text{B} \times 2 = 16$ GB of weights before you serve a single token.

The KV cache is the part that scales with traffic. Every token kept in context stores key and value vectors for each key/value head in every layer:

$M_\text{token} = 2 \times n_\text{layers} \times n_\text{KV heads} \times d_\text{head} \times p_\text{bytes}$

The leading 2 is for the key and the value. For a whole batch, multiply by the batch size $B$ and the sequence length $S$:

$M_\text{KV} = 2 \times B \times S \times n_\text{layers} \times n_\text{KV heads} \times d_\text{head} \times p_\text{bytes}$

Plug in a Llama 3 8B-shaped sketch so the scale is concrete: 32 layers, 32 query heads, head dimension 128, BF16. The first KV line is the multi-head-attention counterfactual (32 key/value heads). Llama 3 8B actually uses grouped-query attention (GQA), sharing 8 key/value heads across those 32 query heads, which is the standard memory lever.^{[1]Reference 1The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783[2]Reference 2GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

At batch 16 and 2K context, this multi-head-attention cache rivals weight memory. Switching from 32 to 8 key/value heads with GQA cuts it from about 17 GB to about 4 GB. First confirm that weights and runtime overhead fit; then KV cache often sets marginal batch and context capacity. Activations and workspace still need a slice of HBM, so don't plan to fill the card to 100%.

A serving prompt asks whether you can raise batch size to improve throughput. Which memory term should you check first, and why?

Answer

After confirming that weights and runtime overhead fit, check KV cache. Weight memory is fixed for a loaded model, while KV memory grows with batch size, sequence length, layers, and key/value heads. Doubling batch or context can exhaust remaining HBM.

Decision log habit

When a design has many possible components, keep a visible decision log:

This keeps the conversation inspectable. Interviewers can disagree with a choice and still see that the choice was deliberate.

45-minute board plan

Practice with a visible clock. Leave the last 10 minutes for follow-ups instead of spending 30 minutes drawing boxes.

If the interviewer interrupts early, jump to the dominant constraint:

The part that most changes the design is constraint. I'll size that first, then show the request path it forces.

Mock design prompts

Treat each prompt like a 45-minute design round. First write requirements, scale math, API, data model, request path, failure modes, and rollout plan. Then open the solution guide.

Prompt 1: model gateway for enterprise teams

Design an API gateway for teams calling multiple LLM providers through one company platform.

Prompt details:

• Each organization has workspaces, users, API keys, and model access rules.
• The gateway must enforce requests/minute, tokens/minute, and monthly spend limits.
• Support needs a request ID that can explain which route, model, policy decision, and overload state happened.
• Some models are beta-only and must be gated.
• Traffic can spike 10x during customer-support incidents.

Clarifying questions to ask:

• Are clients streaming responses, batch jobs, or both?
• Should spend limits be hard stops, soft alerts, or tier-dependent?
• Which decision must support explain first: auth failure, quota failure, route choice, or provider failure?

Solution guide

Put these artifacts on the board:

1. Goal: reliable multi-model access with debuggable controls.
2. API: POST /v1/responses, GET /v1/requests/{id}, admin endpoints for keys and limits.
3. Data model: organization, workspace, key, route policy, model entitlement, usage bucket, request log.
4. Request path: gateway -> auth -> quota estimate -> route policy -> admission queue -> provider adapter -> stream.
5. Overload: return 429 with Retry-After when this tenant, key, or client exceeds its limit. Return 503 with a retry hint when healthy callers can't be admitted because fleet or provider capacity is exhausted. Route to an approved fallback only when product policy allows the behavior change.
6. Observability: request ID, model, token estimate, queue wait, provider latency, error class, policy version.
7. Rollout: canary new routes, kill switch beta models, regression tests for auth and quota bypass.

Common misses: no support path, no versioned policy record, no overload rejection, and no separation between authentication and authorization.

Follow-up guide

If asked about streaming, reserve quota from an estimate before admission, cap output length, then reconcile actual tokens at the end of the stream. If asked about fairness during a 10x spike, split queues by organization or tier so one incident can't starve everyone else. If asked about beta models, answer with a versioned entitlement check and a kill switch:

The gateway should log policy_version, entitlement_id, route_id, and quota_bucket_id on every request so support can explain both accepted and rejected traffic.

Prompt 2: permission-aware enterprise retrieval

Design retrieval for an internal assistant that answers employee questions from company documents.

Prompt details:

• Documents come from multiple systems with different ACL formats.
• Permission changes and deletions must take effect quickly.
• Answers must cite sources.
• The assistant must not retrieve and then filter private documents after generation.
• Admins need auditability for "why did this answer use this document?"

Clarifying questions to ask:

• What is the revocation target: seconds, minutes, or hours?
• Can indexes be physically separated by tenant, or must filters enforce isolation?
• Should the assistant fail closed when ACL freshness is unknown?

Solution guide

Put these artifacts on the board:

1. Ingestion: connector workers pull content plus ACL snapshots and write immutable document versions.
2. Indexing: tenant or workspace isolation plus ACL metadata filters; deleted docs move to a tombstone state.
3. Query path: auth context -> eligible corpus filter -> hybrid retrieval -> rerank -> answer with citations.
4. Freshness: connector lag metric, ACL refresh jobs, deletion queue, and emergency purge path.
5. Audit: query ID, user identity, eligible filters, retrieved doc IDs, citation IDs, policy version.
6. Evals: recall slices by source, permission-denied tests, deletion tests, faithfulness checks.
7. Failure mode: if ACL state is stale or unknown, fail closed for sensitive sources.

Common misses: filtering after generation, no deletion story, no citation IDs, and no way to explain eligibility.

Follow-up guide

If asked about deletion, describe a fast tombstone path first, then slower compaction of embeddings and chunks. If asked about stale permissions, define a freshness SLO per source and fail closed for sensitive sources when the ACL snapshot is too old.

For auditability, store enough to replay the eligibility decision: user identity, groups, source ACL version, query filters, retrieved chunks, citations shown, and model response ID. That lets an admin answer "why this document?" without exposing unrelated private documents.

Prompt 3: long-running coding agent service

Design a service that accepts repository tasks and runs a coding agent asynchronously.

Prompt details:

• Users submit a repo, branch, task prompt, and tool permissions.
• Jobs can run for 30 minutes and may need retries or human review.
• The agent can create commits, run tests, and leave artifacts.
• Users need live progress, cancellation, logs, and final diff review.
• Secrets and production systems must be protected.

Clarifying questions to ask:

• Is the service allowed to push branches, or should it only produce a patch artifact?
• Which tools need network access, and which should run offline?
• What happens when a user cancels during a tool call or test run?

Solution guide

Put these artifacts on the board:

1. API: create task, get task status, stream events, cancel task, list artifacts.
2. State machine: queued, provisioning, running, blocked, review, failed, canceled, complete.
3. Execution boundary: sandbox with scoped repo checkout, network policy, secret redaction, time and disk limits.
4. Persistence: task row, run attempts, tool calls, logs, artifacts, branch/commit refs, cancellation flag.
5. Recovery: checkpoint workspace state, retry transient infra failures, never replay unsafe writes without idempotency.
6. Observability: per-tool latency, test results, token use, queue time, sandbox exits, reviewer actions.
7. Rollout and safety: permission presets, allowlisted tools, audit log, kill switch by tool or model route.

Common misses: no cancellation semantics, no sandbox boundary, no artifact model, and no distinction between retrying a read and replaying a write.

Follow-up guide

If asked about retries, separate infrastructure retries from agent-action retries. Retrying a sandbox provision is safe; replaying a commit, comment, or external API write needs idempotency or human review.

If asked about cancellation, make it cooperative and durable: set a cancellation flag, stop scheduling new tool calls, terminate the sandbox after a grace window, persist partial logs, and mark artifacts as incomplete. If asked about secrets, say the sandbox receives scoped, short-lived credentials and the event log redacts values before storage.

Why is a support/debug path part of the design, not an afterthought?

Answer

Production AI products fail in ways users and operators need to inspect. Request IDs, traces, eligible documents, model choice, policy decisions, and overload state let support debug without guessing.

Drill 1: API gateway and rate-control plane

The gateway is the front door. It authenticates API keys, resolves workspace and organization limits, estimates request cost, routes to a model or queue, and emits a request ID that support can follow. Draw the admit path first. Model routing and streaming hang off this spine; they don't replace it.

Design checklist:

• API keys map to workspace and organization.
• Limits apply across requests/minute, tokens/minute, model, and tier.
• Request IDs appear in responses and logs.
• Tenant or client limits return 429; fleet or provider overload returns 503. Both carry a request ID and a bounded retry hint instead of silent queue growth.
• Beta features are gated by explicit version or feature flags.
• Rollout has canaries, kill switches, and regression checks.

Use Python to sanity-check rate math before drawing capacity boxes:

The estimate is an admission-time reservation, not the final bill. A streaming response needs a maximum output budget, then a reconciliation step when the stream ends:

Drill 2: inference scheduler

For LLM serving, the scheduler is where latency, cost, and fairness meet. NVIDIA documents in-flight batching as a way to interleave context and generation work so GPUs are used more efficiently while latency stays under control.^{[3]Reference 3Paged Attention, IFB, and Request Scheduling.https://nvidia.github.io/TensorRT-LLM/features/paged-attention-ifb-scheduler.html}

Serving physics first

Before any scheduler diagram, separate the two work phases:

That split is why dual SLOs matter: a design can meet mean tokens/sec while still failing stream feel (bad ITL) or first-token UX (bad TTFT). DistServe treats goodput as the request rate that still meets both TTFT and decode-pacing constraints. Colocating the two phases couples those SLOs: a long prefill in the same batch delays every in-flight decode.^{[4]Reference 4DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving.https://arxiv.org/abs/2401.09670} Chunked prefill, separate prefill/decode pools, and prefill/decode disaggregation are answers to this physics, not optional jargon. Chunked prefill reduces interference by bounding each step. Disaggregation removes it by putting the phases on different GPUs.

Metrics to name in the lab loop

• Time to first token (TTFT): arrival → first output token (queue + prefill + first token).
• Inter-token latency (ITL) / time between tokens (TBT): each gap after the first token; use the distribution (including tail) for stream feel.
• Time per output token (TPOT): mean decode pacing after the first token (not a synonym of ITL).
• Queue wait time before admission or first scheduled work.
• p95 and p99 end-to-end latency, plus p95/p99 of TTFT and ITL when the product streams.
• Fleet goodput / admitted RPS under dual TTFT and ITL SLOs (raw tokens/sec without SLO attainment is not capacity).
• Tokens per second only with a label: per-request decode rate, fleet aggregate, or prefill rate. Undifferentiated TPS is not a substitute for ITL or TPOT.
• GPU utilization.
• Error rate and overload rejections.
• Cost per successful request.

Serving toolkit (when to reach for each)

Continuous batching and paged attention are orthogonal: continuous batching decides when sequences join or leave the batch; paged attention decides how KV blocks are laid out and freed. Name both mechanisms separately in an interview.

Model-fit scale math

If weights plus KV don't fit one GPU, say the next move in interconnect language: tensor parallelism (TP) for layer shards over high-bandwidth links (often NVLink within a node), pipeline parallelism (PP) across stages (more tolerant of inter-node links, pays bubble cost), and expert parallelism when MoE routing requires it. Always state the communication cost you accept for the fit.

Start capacity planning with a measured workload-specific fleet capacity, not a generic GPU estimate:

For an interview-sized overload policy, estimate queue wait before admission. This approximation deliberately stays simple: it treats tokens as homogeneous. A real scheduler tracks two resources separately (prefill token backlog vs active decode sequences / free KV blocks) and measured tail latency for each SLO.

Keep rejection ownership visible. 429 Too Many Requests says the caller exceeded a tenant, key, or client policy and could succeed after its quota window resets. 503 Service Unavailable says the service fleet or an upstream provider lacks capacity for an otherwise eligible request. A full tenant bucket can produce 429 even when GPUs are idle; a saturated fleet can produce 503 for a tenant that is under quota.

The wait check below is the 503 side of that split. 20,000 queued tokens clear a 0.25 s budget at 125,000 tokens/s. 50,000 don't. The two-resource sketch then adds decode slots and free KV blocks so a request can fail admission even when the prefill backlog looks fine.

When should the system return 429 or 503 instead of queueing more work?

Answer

Reject before queueing violates latency or fairness. Use 429 when the caller's own quota or rate policy blocks admission. Use 503 when the shared fleet or provider is overloaded despite the caller remaining eligible. Include Retry-After when the service can estimate a safe retry window.

Drill 3: permission-aware retrieval

For enterprise retrieval, the critical rule is: don't retrieve private data and filter it after generation. Permission constraints must be part of candidate selection, ranking, and auditing.

Architecture pieces:

• Connector ingestion workers with backpressure.
• Per-document ACLs or delegated auth checks.
• Tenant-isolated indexes or strict metadata filters.
• Hybrid retrieval plus reranking.
• Citation output with source IDs.
• Deletion and retention jobs.
• Offline evals for recall and answer faithfulness.
• Support traces that show which documents were eligible.

The ACL snapshot and hybrid index both feed candidate selection. That connection matters: unauthorized chunks shouldn't reach the reranker or model context.

For sensitive sources, encode the freshness decision as a fail-closed policy:

Filter tombstones and ACLs before ranking. The tiny fixture below makes the order visible:

Drill 4: long-running coding agents

Long-running agent infrastructure has to persist intent, tool calls, artifacts, checkpoints, logs, and permissions. The main design risk isn't just failed execution. It's uncontrolled execution.

Cover:

• Task states: queued, provisioning, running, blocked, review, failed, canceled, complete.
• Checkpoints for resumability.
• Tool permission scopes and audit logs.
• Secret redaction.
• Git branch and conflict handling.
• Streaming progress.
• Cancellation and deadlines.
• Evals for task success and regression.

Write legal transitions down before discussing workers. That prevents a cancellation or retry from jumping into an impossible state:

Retries need a policy boundary too. Reads and sandbox provisioning can retry automatically. External writes need an idempotency key or review:

Drill 5: eval gates and staged rollout

An evaluation monitor isn't just a dashboard. It turns fixed expectations and incident discoveries into a regression suite, then blocks launch when a critical slice fails.

Keep online canaries reversible. Offline evals can pass while latency, provider errors, or permission failures regress under real traffic:

Common serving misconceptions

For each row, name the old approach's failure and the mechanism that replaces it. Don't stop at "we use vLLM now."

Failure modes to avoid

• Starting with implementation technology before naming user value.
• Caching without invalidation, privacy, or freshness.
• Monitoring without exact metrics.
• Ignoring overload and support/debug needs.
• Treating safety as a slogan instead of evals, permissions, staged rollout, and rollback.
• Forgetting that agent systems need reversible actions and audit trails.