AI Lab System Design Interview

A system-design round asks you to turn ambiguous model-product requirements into a reliable backend. A strong answer is rarely the most complex architecture. It names the product goal, sizes the hard constraint, then adds queues, caches, model routing, eval gates, permissions, or human review only where requirements force them.

Use codebase, model-serving, RAG, support-automation, and deployment examples when you need a concrete neutral domain: code search, incident summarization, policy-gated billing credits, eval dashboards, and internal agent workflows have real latency, privacy, and debugging constraints.

The design round script

Open with:

I will 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:

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.

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 will 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 teach the interviewer 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 8-billion-parameter 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 simple numbers so the scale is concrete. The GQA line shows why grouped-query attention (sharing key/value heads across query heads) is the standard memory lever:

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.

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. A strong design round leaves time 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 will 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

Strong answer shape:

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

Strong answer shape:

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

Strong answer shape:

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.

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.^{[1]Reference 1Paged 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). Chunked prefill, separate prefill/decode pools, and prefill/decode disaggregation are answers to this physics, not optional jargon.

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.

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 correction, identify old approach's failure and replacement mechanism.

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.