How to Prepare for an AI Infrastructure Engineer Interview

The AI infrastructure engineer interview is the most technically demanding screen in the current job market. The role sits at the intersection of hardware, distributed systems, and deep learning — and the people hiring for it know exactly what they're looking for.

If your background is web services, you'll need to learn a new vocabulary: warps, FLOPS, tensor cores, NVLink, AllReduce. If your background is ML research, you'll need to learn engineering discipline: latency budgets, fault tolerance, operational runbooks. This guide gives you both sides — what the role actually involves, what interviewers ask, and how to prepare.

---

What AI Infra Engineers Actually Build

AI infrastructure engineers build the systems that make large-scale AI work. That means different things at different companies:

At AI labs (Anthropic, OpenAI, DeepMind, xAI): You're building and operating GPU clusters — sometimes thousands of GPUs — for pre-training large models. You write distributed training frameworks, optimize collective communication (AllReduce, AllGather), and debug training runs that cost $10M+ and can't just be restarted from scratch.

At hyperscalers (Meta, Microsoft, AWS): You're building inference infrastructure that serves billions of requests per day. Latency is everything. You're optimizing CUDA kernels, deploying custom hardware (AWS Trainium, Microsoft Maia, Meta MTIA), and building serving systems that maintain < 100ms p99 latency at enormous scale.

At ML-heavy companies (Netflix, Spotify, Uber, Airbnb): You're building the training and serving infrastructure used by internal ML teams. Less about cutting-edge hardware, more about reliability and developer productivity — making it easy for 50 ML engineers to train and deploy models without fighting infrastructure.

---

Who's Hiring and the Salary Reality

This is a small talent pool. The number of engineers who genuinely understand GPU architecture, distributed training, and inference optimization is limited — which is why compensation is exceptional.

Companies actively hiring: NVIDIA, Meta AI Research (PyTorch infra team), Microsoft AI (Azure AI infra), Cerebras Systems, Groq, SambaNova Systems, Amazon (Trainium/Inferentia teams), Mistral, and any funded AI startup that's training models.

The honest truth: if you don't have GPU programming experience (CUDA or Triton), you won't pass the technical screen at an AI lab. If your target is a hyperscaler's AI infrastructure team, strong distributed systems knowledge plus ML serving experience is more important than CUDA depth.

---

Interview Structure

At AI labs, expect a CUDA or Triton kernel optimization question. At hyperscalers, expect heavier systems design. At ML-heavy companies, expect more MLOps-adjacent questions about serving and infrastructure tooling.

---

Topic 1: GPU Architecture Fundamentals

You need to understand why GPUs are designed the way they are before you can optimize for them.

The key architectural difference from CPUs:

A CPU has a few powerful cores (8–64) optimized for sequential tasks with complex control flow. A GPU has thousands of simpler cores organized into Streaming Multiprocessors (SMs), optimized for parallel tasks with minimal branching. An NVIDIA H100 has 132 SMs, each with 128 CUDA cores — that's 16,896 CUDA cores total.

The CUDA execution model:

Threads are grouped into warps of 32 threads. All threads in a warp execute the same instruction simultaneously (SIMT — Single Instruction, Multiple Threads). If threads in the same warp take different branches of an if statement, all branches are executed and the wrong ones are masked. This is warp divergence and it kills performance.

Triton: OpenAI released Triton in 2021 as a Python-like language for writing GPU kernels without deep CUDA expertise. Instead of thinking in threads and warps, you think in tiles (blocks of elements). PyTorch 2.0's torch.compile uses Triton under the hood to auto-generate fused kernels. For interviews at NVIDIA, Meta (PyTorch team), or any company working on kernel optimization, Triton knowledge is increasingly expected alongside CUDA. FlashAttention 2 is implemented in Triton.

Memory hierarchy (fastest to slowest):

| Memory | Size | Latency | Bandwidth | |--------|------|---------|-----------| | Registers | 256KB/SM | ~1 cycle | — | | Shared memory (L1 cache) | 228KB/SM | ~32 cycles | ~20TB/s | | L2 cache | 50MB (H100) | ~200 cycles | ~7TB/s | | HBM (global memory) | 80GB (H100) | ~600 cycles | 3.35TB/s |

The bandwidth wall: Most deep learning operations are memory-bandwidth-bound, not compute-bound. An H100 has 1979 TFLOPS of BF16 compute but only 3.35TB/s of memory bandwidth. A simple matrix multiply of large matrices can be compute-bound; an elementwise operation (like a GELU activation) is almost always memory-bandwidth-bound.

This is why kernel fusion matters so much: instead of reading from HBM, applying GELU, writing back to HBM, then reading again for the next operation — you fuse the operations into one kernel that keeps data in shared memory across multiple operations. FlashAttention works this way.

H100 vs A100:

---

Topic 2: Distributed Training

No single GPU can train a 70B parameter model — the weights alone require 140GB in BF16, and you need optimizer states and gradients on top of that. Distributed training splits the work across many GPUs.

Three parallelism strategies:

Data Parallelism (DP): Replicate the full model on each GPU. Split the batch across GPUs. Each GPU computes gradients on its shard, then gradients are averaged across all GPUs (AllReduce). Simple and efficient for models that fit on a single GPU. PyTorch's DistributedDataParallel (DDP) implements this.

Tensor Parallelism (TP): Split individual weight matrices across GPUs. For a large linear layer, GPU 0 holds the first half of the weight matrix, GPU 1 holds the second half. Both compute in parallel, then results are concatenated (AllGather). Used within a single node because it requires high-bandwidth interconnects (NVLink). Megatron-LM pioneered this.

Pipeline Parallelism (PP): Split the model's layers across GPUs. GPU 0 runs layers 1–10, GPU 1 runs layers 11–20, etc. A micro-batch flows through the pipeline stage by stage. The challenge: GPU 0 is idle while GPU 1 processes. Addressed by micro-batching (GPipe) and interleaved schedules (PipeDream).

In practice, large-scale training uses 3D parallelism: tensor + pipeline + data parallelism combined. A 70B model training on 1024 GPUs might use 8-way TP within a node, 4-way PP across nodes, and 32-way DP across pipeline groups.

ZeRO (Zero Redundancy Optimizer):

Standard DDP replicates the full optimizer state on each GPU. For a large model, the optimizer state (Adam: 2× model params in FP32) dominates memory. ZeRO partitions this across GPUs:

• ZeRO-1: partition optimizer states
• ZeRO-2: + partition gradients
• ZeRO-3: + partition model parameters

ZeRO-3 effectively removes the memory redundancy of DP while keeping the simplicity. DeepSpeed implements ZeRO.

FSDP (Fully Sharded Data Parallel): PyTorch's built-in alternative to DeepSpeed ZeRO-3, available from PyTorch 1.12. FSDP shards model parameters, gradients, and optimizer states across all ranks — equivalent to ZeRO-3. The key difference: FSDP is native PyTorch with no external dependencies, making it the default choice for new projects unless you need DeepSpeed-specific features (like ZeRO-Infinity for CPU offloading). Most production training at Meta, Microsoft, and Hugging Face now uses FSDP.

Gradient checkpointing (activation checkpointing): During the forward pass, intermediate activations are stored in GPU memory for use in the backward pass — this can consume 3–5× more memory than the model weights alone for long sequences. Gradient checkpointing trades compute for memory: instead of storing all activations, only store activations at checkpoint boundaries and recompute intermediate activations during the backward pass. Enables training with roughly 10× longer sequences at the cost of ~33% more compute. Enable with model.gradient_checkpointing_enable() in Hugging Face Transformers.

Collective operations you must know:

• AllReduce: sum gradients across all ranks — each rank ends up with the same sum. Used in DDP.
• AllGather: collect tensors from all ranks — each rank ends up with the concatenated tensor. Used in tensor parallelism.
• ReduceScatter: reduce and scatter — each rank gets a different shard of the reduced result. Used in ZeRO and ring-AllReduce.

---

Topic 3: Inference Optimization

Training is expensive but offline. Inference is cheaper per run but must happen at low latency, at massive scale, continuously. The optimization goals are different.

KV Cache:

During autoregressive generation, the attention mechanism needs the Key and Value vectors for all previous tokens. Without caching, you'd recompute them for every new token — O(n²) complexity per token generated.

KV cache: store the K and V tensors from all previous positions in GPU memory. For each new token, only compute K and V for the new position, then attend over the cached + new K/V. This reduces the compute from O(n²) to O(n) per token.

Cost: KV cache grows with sequence length and batch size. For a 70B model with 4096-token context and batch size 64: KV cache = 2 × num_layers × seq_len × num_heads × head_dim × dtype_bytes. This can be tens of GB — often the bottleneck for serving.

PagedAttention (vLLM):

The key insight from the vLLM paper (Kwon et al., 2023): KV caches have variable size (different requests have different sequence lengths) and must be pre-allocated. Traditional serving systems waste 60–80% of GPU memory on fragmentation and over-allocation.

PagedAttention manages the KV cache like an OS manages virtual memory. The cache is divided into fixed-size pages. Pages are allocated on demand and can be non-contiguous in physical memory. This enables sharing pages across requests with the same prompt prefix (common in system-prompt-heavy workloads), reducing memory usage dramatically.

Result: vLLM achieves 2–4× higher throughput than Hugging Face TGI on the same hardware.

Quantization:

Reduce the precision of model weights to save memory and speed up compute.

• FP16/BF16: Standard — no quality loss, 2× memory vs FP32
• INT8: 2× further reduction. Minor quality degradation on some models. Works well with bitsandbytes library.
• INT4/NF4: 4× reduction vs FP16. Noticeable quality loss on some tasks. Used in QLoRA for fine-tuning, and GGUF format for CPU inference. GPTQ and AWQ are the most common post-training quantization methods.
• FP8 (H100-native): H100 Transformer Engine supports FP8 training and inference with minimal quality loss — 2× throughput vs BF16.

Speculative decoding:

Standard autoregressive decoding generates one token per forward pass — expensive because each pass runs the full model. Speculative decoding uses a small draft model to generate multiple candidate tokens cheaply, then verifies them with the full model in a single forward pass.

If the draft model's tokens match what the full model would have generated, they're accepted. If not, reject and fall back. The speedup depends on the acceptance rate — typically 2–3× for matched draft/target model pairs (e.g., GPT-2 as draft, GPT-4 as target).

---

Topic 4: Networking for AI

Inter-GPU communication is often the bottleneck in distributed training, not compute. Knowing this layer separates strong candidates.

NVLink: NVIDIA's proprietary high-bandwidth interconnect between GPUs within a single server. H100 NVLink bandwidth: 900GB/s bidirectional (vs PCIe 5.0 at ~128GB/s). NVLink is why tensor parallelism works well within a single node — you have enough bandwidth to pass partial activations frequently.

InfiniBand: The standard for inter-node communication in GPU clusters. HDR InfiniBand: 200Gb/s (25GB/s). NDR InfiniBand: 400Gb/s. AI supercomputers like the NVIDIA DGX SuperPOD use InfiniBand for the scale-out network.

RDMA (Remote Direct Memory Access): A protocol that allows a GPU on one server to read/write another server's GPU memory directly, bypassing the CPU. Critical for low-latency collective operations. GPUDirect RDMA enables this.

Interview question: "Why does tensor parallelism work well within a node but poorly across nodes?"

Answer: Tensor parallelism requires frequent AllGather operations to synchronize partial activations. These generate a lot of small, latency-sensitive messages. NVLink (within node) has ~50ns latency and 900GB/s bandwidth. InfiniBand (across nodes) has ~1µs latency and 25GB/s bandwidth. The 20× bandwidth difference and 20× latency difference make frequent cross-node AllGather expensive enough to dominate the computation time.

---

10 Interview Questions with Answers

Q1: What is the memory bandwidth bottleneck in LLM inference?

Most LLM inference operations (especially the attention mechanism and linear layers during decode) are memory-bandwidth-bound, not compute-bound. The GPU can execute arithmetic faster than it can feed data from HBM. The roofline model helps: operations below the arithmetic intensity ridge are bandwidth-bound. Solution: quantization (fewer bytes moved per weight), kernel fusion (reduce the number of HBM read/write cycles), and larger batch sizes (amortize the memory reads across more computations).

Q2: Explain how KV cache works and what its limitations are.

KV cache stores the Key and Value projections for all previous tokens so they don't need to be recomputed at each generation step. Limitation 1: memory — cache grows linearly with sequence length × batch size × num_layers × head_dim. Limitation 2: pre-allocation — traditional systems pre-allocate the max sequence length, causing fragmentation. PagedAttention (vLLM) solves limitation 2. Limitation 3: multi-turn conversations — cache must be evicted when context exceeds the limit, causing quality degradation.

Q3: What is warp divergence and how do you avoid it?

Warp divergence occurs when threads in the same warp take different code paths (if-else branches). Since all threads must execute the same instruction, divergent threads are masked and serialized — reducing effective parallelism. Avoid it by restructuring code to ensure uniform control flow within a warp, or by padding and sorting inputs so similar-length inputs are grouped in the same warp (common in NLP batching).

Q4: How would you implement data-parallel distributed training in PyTorch?

Use DistributedDataParallel (DDP). Initialize a process group with torch.distributed.init_process_group("nccl"). Wrap your model with DDP(model). Use a DistributedSampler so each process gets a non-overlapping data shard. DDP automatically registers hooks to AllReduce gradients after the backward pass before the optimizer step. Launch with torchrun which handles rank assignment.

Q5: What is the difference between tensor parallelism and pipeline parallelism?

Tensor parallelism splits individual operators (weight matrices) across GPUs — all GPUs work on the same layer simultaneously and synchronize results. Requires high bandwidth (NVLink). Good for layers with large weight matrices. Pipeline parallelism splits different layers across GPUs — each GPU runs a different stage of the network. Lower bandwidth requirement but introduces pipeline bubbles (idle time waiting for previous stages). They're complementary and often used together.

Q6: Explain ZeRO-3.

ZeRO-3 partitions optimizer states, gradients, AND model parameters across data-parallel ranks. No single GPU holds the complete model. During the forward pass, parameters are gathered (AllGather) when needed by each layer, then immediately discarded. During backward, parameters are re-gathered, gradients are computed, then ReduceScattered. This enables training models larger than a single GPU's memory while maintaining data parallelism's simplicity.

Q7: How does speculative decoding achieve speedup?

A small draft model generates K candidate tokens quickly and cheaply. The large verifier model processes all K candidates in parallel in a single forward pass (vs K sequential forward passes). Tokens that match the verifier's distribution are accepted; the first rejected token is replaced by the verifier's sample. The expected accepted tokens per verifier call depends on the acceptance rate — if the draft model is accurate 75% of the time, you get ~2.5× speedup with K=8 drafts.

Q8: Design a serving system for a 70B parameter model serving 10,000 concurrent users.

Start with capacity estimation: 70B model in BF16 = 140GB. An H100 has 80GB — need at least 2 H100s just for weights. Add KV cache for serving: at 4096 context, batch 64, ~30GB more. So 4 H100s per replica. For 10K concurrent users at ~2 requests/second each = 20K RPS. With continuous batching and throughput of ~500 tokens/sec per replica at batch 64, estimate and size accordingly. Use vLLM for serving, a load balancer with consistent hashing, and horizontal scaling with multiple replicas.

Q9: What is Flash Attention and why is it faster?

Standard attention computes the full attention matrix (N×N for sequence length N) and stores it in HBM — O(N²) memory and multiple HBM read/write cycles. FlashAttention (Dao et al., 2022) tiles the attention computation: it processes blocks of Q, K, V in SRAM (shared memory) and uses a numerically stable online softmax to compute the correct output without ever materializing the full N×N matrix in HBM. Result: memory-bandwidth operations reduced by 5–20×, significant speedup for long sequences.

Q10: How do you debug a distributed training run that's significantly slower than expected?

Profile with NVIDIA Nsight Systems to get a GPU timeline. Look for: (1) low GPU utilization — might be a data loading bottleneck (fix: increase DataLoader workers, use prefetching); (2) long AllReduce times — might be a slow node in the cluster (use NCCL's debug mode to identify); (3) pipeline bubbles — if using pipeline parallelism, check the ratio of compute to communication. Use torch.profiler with Chrome trace visualization to see operator-level timing.

---

4-Month Preparation Plan

---

Practice on KnowEaseAI

AI infrastructure is the discipline where most self-study resources either go too shallow (blog posts that explain GPU basics at surface level) or too deep (academic papers that assume prior CUDA expertise). The gap is practical interview preparation.