AI infrastructure engineer domain knowledge interview questions

As asked

What is the time and memory complexity of vanilla self-attention with respect to sequence length N and model dimension D? Why does this create a problem at long contexts, and name two architectural approaches that reduce the attention complexity.

Sample answer outline

Time complexity is O(N^2 * D) and memory is O(N^2) to store the attention matrix. At N=128K tokens the N^2 term becomes prohibitive. Approaches: sparse attention (BigBird, Longformer) limits each token to attend to a subset of positions; linear attention (Performer, RWKV kernel trick) approximates softmax attention in O(N*D). FlashAttention is not a complexity reduction but reduces IO complexity. Strong candidates can distinguish IO complexity from computational complexity.

Expect these follow-ups

• Does FlashAttention change the asymptotic computational complexity of attention, or only the IO complexity?

As asked

How much GPU memory does the Adam optimizer consume for a 7B parameter model in mixed-precision training, and why? Break it down by component (parameters, gradients, optimizer state).

Sample answer outline

Parameters: 7B * 2 bytes (bf16) = 14 GB. Gradients: 7B * 2 bytes = 14 GB (can be kept in bf16 but often accumulated in fp32). Optimizer state for Adam: fp32 master weights (7B * 4 = 28 GB), fp32 first moment (28 GB), fp32 second moment (28 GB). Total optimizer state = 84 GB. Grand total before activations is approximately 112 GB to 126 GB. This is why ZeRO and FSDP are necessary: a single 80 GB A100 cannot hold a 7B model with Adam state.

Expect these follow-ups

• How does AdamW differ from Adam in terms of memory footprint, and why did AdaFactor see adoption for large model training?

As asked

Describe the GPU memory hierarchy from registers to HBM on an A100. Give approximate sizes and bandwidths for each level, and explain how a CUDA kernel programmer should think about data placement across these levels.

Sample answer outline

Registers: per-thread, fastest, ~64K 32-bit registers per SM. Shared memory (L1): 192 KB per SM on A100, bandwidth in TB/s range, shared within a thread block. L2: 40 MB on A100, hundreds of GB/s. HBM: 80 GB, 2 TB/s. The programmer should maximize register and shared memory reuse (keep hot data on-chip), minimize HBM accesses, and avoid bank conflicts in shared memory. Mention that shared memory is explicitly managed by the programmer unlike CPU caches.

Expect these follow-ups

• What is the maximum amount of shared memory per thread block on A100, and how does using more shared memory per block affect occupancy?

As asked

Explain NUMA topology and why it matters for GPU training servers. What happens when a GPU and its associated HCA are on different NUMA nodes, and how do you verify and fix NUMA affinity for a training process?

Sample answer outline

NUMA (Non-Uniform Memory Access) splits CPU cores and memory into nodes; cross-NUMA access adds latency and reduces bandwidth. In an 8-GPU server, GPUs 0-3 are typically on NUMA node 0 and GPUs 4-7 on NUMA node 1. If a GPU's IB HCA is on a different NUMA node, data for RDMA operations crosses the UPI/QPI interconnect, adding latency and reducing bandwidth. Verify with numactl --hardware and nvidia-smi topo -m. Fix by pinning the training process to the correct NUMA node with numactl --cpunodebind=X --membind=X or setting CPU affinity in the job launcher.

Expect these follow-ups

• How does numactl interact with NCCL's GPU-HCA affinity selection, and can NCCL detect and correct NUMA misalignment automatically?

As asked

Explain why transformer LLMs use LayerNorm rather than BatchNorm, particularly in the context of distributed training. What collective communication would BatchNorm require in a data-parallel setup?

Sample answer outline

BatchNorm computes statistics over the batch dimension, which in data-parallel training is split across GPUs. Correct batch statistics would require an allreduce of sum and sum-of-squares across all ranks per forward pass, adding significant communication overhead and latency. LayerNorm operates over the feature dimension (hidden dim) within each sample independently, so no inter-GPU communication is needed. Additionally, BatchNorm behaves poorly with very small per-GPU batch sizes common in LLM training.

Expect these follow-ups

• What is RMSNorm and why is it preferred over LayerNorm in some newer LLM architectures like Llama?

As asked

What are rotary position embeddings (RoPE) and how do they encode positional information differently from learned absolute position embeddings? What does RoPE extrapolation to longer sequences require at inference time?

Sample answer outline

RoPE applies a rotation matrix to the query and key vectors based on their position, such that the dot product between q at position i and k at position j depends only on their relative offset i-j. Unlike absolute embeddings that add a learned vector, RoPE has no learned parameters and naturally encodes relative distance. Extrapolation beyond training context requires scaling the rotation frequencies (YaRN, LongRoPE, position interpolation) because the model has never seen the large rotation angles corresponding to long-range positions.

Expect these follow-ups

• How does YaRN extend a model's context length without full fine-tuning, and what is the core mathematical trick?

As asked

What does torch.backends.cudnn.benchmark = True do, and when should you not enable it? What is the trade-off between enabling it and reproducibility of results?

Sample answer outline

cudnn.benchmark runs a set of cuDNN algorithm candidates for each convolution shape on the first call and caches the fastest one. This improves throughput when input shapes are fixed but wastes time at startup and can introduce non-determinism because different cuDNN algorithms may round floating-point operations differently. Should not be enabled when input shapes vary (causes repeated benchmarking) or when exact reproducibility is required. For transformers with fixed sequence lengths it is generally safe to enable.

Expect these follow-ups

• How do you achieve deterministic results in PyTorch training when cudnn.benchmark is off, and what other sources of non-determinism remain?

As asked

What matrix dimension and data type requirements must a GEMM operation satisfy to utilize Tensor Cores on A100? What happens to throughput when these requirements are not met, and how would you pad a model to satisfy them?

Sample answer outline

On A100, Tensor Cores require matrix dimensions (M, N, K) to be multiples of 8 at minimum for fp16 and bf16; in practice padding to multiples of 64 or 128 gives the best throughput by aligning to the warp tile granularity. tf32 has the same alignment requirements as fp16 on A100 since the underlying warp tile is 16x16x8. If dimensions are not aligned, CUDA falls back to scalar CUDA cores with much lower throughput. Padding hidden dimensions to the next multiple of 64 or 128 is common. In practice, model hidden sizes like 4096, 8192, and 12288 all satisfy this. Vocabulary size is a common culprit: an odd vocab size will cause the output projection GEMM to fall off Tensor Cores.

Expect these follow-ups

• How would you check whether a specific GEMM in your model is using Tensor Cores vs. scalar cores in a Nsight Compute profile?

As asked

What is FP8 and what are the two FP8 formats supported on H100? Why does FP8 training require scaling factors, and what makes it more complex to implement correctly than bf16 training?

Sample answer outline

FP8 has two variants on H100: E4M3 (4 exponent, 3 mantissa bits, higher precision but smaller dynamic range) and E5M2 (5 exponent, 2 mantissa bits, larger dynamic range but lower precision, typically used for gradients). FP8 has very limited dynamic range, so per-tensor or per-block scaling factors must be maintained to keep values in range. This requires delayed scaling (using statistics from previous steps) or just-in-time scaling. The main complexity is that both forward activations and gradients need calibrated scales, and the scale must be communicated across distributed ranks for allreduce correctness.

Expect these follow-ups

• How does Transformer Engine from NVIDIA automate the scaling factor management for FP8 training?

As asked

Explain the RDMA programming model: what is a queue pair, what is the difference between send/receive and RDMA read/write operations, and why is RDMA read expensive compared to RDMA write for latency-sensitive workloads?

Sample answer outline

A queue pair (QP) consists of a send queue and a receive queue managed by the HCA. RDMA write/read allow one side to directly access the other's registered memory without involving the remote CPU. RDMA write posts a work request on the local send queue; the HCA executes the DMA without remote CPU involvement. RDMA read requires the local HCA to issue a read to the remote HCA and wait for data return, making it a round-trip operation vs. the one-way write. For low-latency messaging, RDMA write + doorbell notification is preferred over RDMA read.

Expect these follow-ups

• What is the difference between RC (Reliable Connected) and UD (Unreliable Datagram) transport in InfiniBand, and which does NCCL primarily use?