AI infrastructure engineer coding interview questions

Q: Write a Python script that samples nvidia-smi every 5 seconds and logs per-GPU utilization and memory usage to a CSV, with timestamps. The script should handle the case where a GPU disappears from the listing (e.g., after a driver reset).

The script should use subprocess to run nvidia-smi --query-gpu and --format=csv with a loop, parse the CSV output, handle missing or empty GPU rows gracefully, and write a timestamped CSV. A strong answer uses nvidia-smi's built-in --loop flag or subprocess.Popen for streaming, handles subprocess.CalledProcessError on driver failure, and notes that pynvml is a more reliable Python interface than parsing text output.

Q: Write a function that takes a model config (num_layers, num_heads, head_dim, dtype_bytes) and serving config (max_seq_len, max_batch_size) and returns the peak KV cache memory in bytes. Then compute it for Llama-2-70B in fp16 serving 32 concurrent requests at 4096 tokens each.

KV cache = 2 * num_layers * num_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size. For Llama-2-70B: 80 layers, 64 heads (GQA uses 8 KV heads), head_dim=128, fp16=2 bytes, 4096 tokens, 32 requests. With GQA the answer must use num_kv_heads (8) not num_heads. This is 2 * 80 * 8 * 128 * 2 * 4096 * 32 = approximately 42 GB. A strong candidate mentions GQA/MQA and its memory savings.

Q: You have N=4 ranks. Ranks 0, 1, and 3 call dist.all_reduce(), but rank 2 mistakenly calls dist.broadcast() with the same tensor. Write a minimal reproducer and describe what you observe and how you fix it.

The collective type mismatch causes ranks to send incompatible protocol headers to each other; NCCL will either hang indefinitely or raise an internal error after a timeout. The reproducer should launch 4 processes with torch.multiprocessing or torchrun, have rank 2 call a different collective, and demonstrate the hang. Fix by ensuring all ranks call the same collective with the same arguments. Defensive pattern: use TORCH_DISTRIBUTED_DEBUG=DETAIL to log collective call sites.

Q: When using gradient accumulation with PyTorch DDP and accumulating over K micro-batches before each optimizer step, what is the incorrect pattern that causes unnecessary allreduce calls on every micro-batch, and how do you fix it?

DDP triggers allreduce on every backward() call by default. Without no_sync(), every micro-batch backward triggers an allreduce, wasting communication for intermediate accumulation steps. Fix: wrap all micro-batch backward passes except the last in model.no_sync() context manager. The last backward triggers the real allreduce. Also need to scale gradients by 1/K or loss by 1/K before accumulation to match full-batch semantics.

Q: Write a simulation that models how increasing batch size affects per-token latency and throughput for an LLM decode step, assuming decode is memory-bandwidth-bound. Show how you would determine the optimal batch size given a latency SLA of 50 ms per token.

In the memory-bandwidth-bound regime, decode time per step scales linearly with batch size (each sequence reads the full model weights once). Throughput = batch_size / latency_per_step, which is flat in the bandwidth-bound regime since latency also scales linearly. The model should show that throughput saturates once compute-bound, and the optimal batch size for a 50 ms SLA is bandwidth / (model_bytes / 0.05). A strong answer distinguishes prefill (compute-bound) from decode (memory-bound).

Q: Implement a simple ring-allreduce for a list of tensors across N ranks. Each rank should send its chunk to the next rank and receive from the previous rank, accumulating partial reductions. Show the reduce-scatter phase only.

The implementation should split the tensor into N chunks, iterate N-1 times sending the current chunk to rank+1 and receiving from rank-1, accumulating the received chunk into the local buffer. Key correctness points: non-blocking sends/receives to avoid deadlock, correct wraparound for rank 0 and rank N-1, and proper chunk indexing. A strong candidate notes that a real implementation uses CUDA-aware MPI or NCCL and avoids this Python approach in production.

Q: Write a function that takes pipeline depth P and number of micro-batches M and computes the bubble fraction (fraction of time the pipeline is not fully occupied) using the 1F1B (one forward one backward) schedule.

Bubble fraction = (P - 1) / M for the 1F1B schedule. The function should verify M >= P for a valid schedule, compute the formula, and optionally visualize the schedule as a grid. A strong candidate knows this comes from the Megatron paper and explains why increasing M reduces the bubble at the cost of more micro-batch overhead.

10 questions on coding for ai infrastructure engineer candidates. Expect prompts such as “Allreduce communication volume calculation” and “Roofline model for CUDA kernel analysis”, each with a worked answer outline and the follow-ups interviewers push on.

Showing 1 to 10 of 10 coding questions.

As asked

You are training a model with 7B parameters using DDP across 64 GPUs. Calculate the per-GPU data transmitted during one allreduce of the full gradient, assuming bf16 precision. How does this change if you switch to FSDP FULL_SHARD?

Sample answer outline

DDP allreduce: each GPU sends and receives 2 * (N-1)/N * message_size; for N=64 this is approximately 2 * message_size. 7B params * 2 bytes (bf16) = 14 GB per allreduce. With FSDP FULL_SHARD, the reduce-scatter transmits message_size * (N-1)/N and all-gather transmits the same, totaling 2 * 14 GB / 64 * (N-1) roughly but each rank only holds 1/N of the parameters, so peak in-flight per rank is much smaller. The candidate should derive the formula not just give the number.

Reference implementation (python)

Python

# Estimate allreduce traffic for DDP vs FSDP
num_params = 7e9
bytes_per_param_bf16 = 2  # bf16
world_size = 64

# DDP: ring-allreduce sends 2*(N-1)/N * data per GPU
ddp_bytes_per_gpu = 2 * (world_size - 1) / world_size * num_params * bytes_per_param_bf16
print(f"DDP allreduce per GPU: {ddp_bytes_per_gpu / 1e9:.2f} GB")

# FSDP FULL_SHARD: reduce-scatter + all-gather, each ~data/N per rank per step
# Each rank stores 1/N params, but communicates full gradient volume in sharded fashion
fsdp_reduce_scatter = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size
fsdp_all_gather = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size
print(f"FSDP reduce-scatter per GPU: {fsdp_reduce_scatter / 1e9:.2f} GB")

Expect these follow-ups

How would you measure actual allreduce bandwidth vs. theoretical peak bandwidth in your cluster?

allreducecommunicationbandwidthfsdpddp

As asked

Given a kernel that performs 512 GFLOPS of compute and reads/writes 8 GB of data, is it compute-bound or memory-bandwidth-bound on an A100 SXM? Show your calculation and explain what the roofline model predicts for the optimal kernel runtime.

Sample answer outline

A100 SXM peak bf16 TFLOPS is 312 TFLOPS; peak HBM2e bandwidth is 2 TB/s. Arithmetic intensity = 512 GFLOPS / 8 GB = 64 FLOPS/byte. Ridge point = 312e12 / 2e12 = 156 FLOPS/byte. Since 64 < 156, the kernel is memory-bandwidth-bound. Roofline-predicted time = 8 GB / 2 TB/s = 4 ms. A strong candidate should also discuss how to improve arithmetic intensity via kernel fusion.

Reference implementation (python)

Python

# Roofline model calculation
peak_tflops = 312e12  # A100 SXM bf16
peak_bandwidth = 2e12  # bytes/s, A100 HBM3

kernel_flops = 512e9
kernel_bytes = 8e9

arithmetic_intensity = kernel_flops / kernel_bytes  # FLOPS/byte
ridge_point = peak_tflops / peak_bandwidth

if arithmetic_intensity < ridge_point:
    bound = "memory-bandwidth"
    roofline_time = kernel_bytes / peak_bandwidth
else:
    bound = "compute"
    roofline_time = kernel_flops / peak_tflops

print(f"Arithmetic intensity: {arithmetic_intensity:.1f} FLOPS/byte")
print(f"Ridge point: {ridge_point:.1f} FLOPS/byte")
print(f"Bound: {bound}")
print(f"Roofline predicted time: {roofline_time*1000:.2f} ms")

Expect these follow-ups

If you fused two consecutive elementwise operations on this kernel, how would the arithmetic intensity change?

rooflinecudamemory-bandwidthperformancea100

As asked

Write a Python script that samples nvidia-smi every 5 seconds and logs per-GPU utilization and memory usage to a CSV, with timestamps. The script should handle the case where a GPU disappears from the listing (e.g., after a driver reset).

Sample answer outline

The script should use subprocess to run nvidia-smi --query-gpu and --format=csv with a loop, parse the CSV output, handle missing or empty GPU rows gracefully, and write a timestamped CSV. A strong answer uses nvidia-smi's built-in --loop flag or subprocess.Popen for streaming, handles subprocess.CalledProcessError on driver failure, and notes that pynvml is a more reliable Python interface than parsing text output.

Reference implementation (python)

Python

import subprocess, csv, time, datetime, sys

FIELDS = ["index", "utilization.gpu", "memory.used", "memory.total"]

def sample_gpus():
    try:
        result = subprocess.run(
            ["nvidia-smi", f"--query-gpu={','.join(FIELDS)}",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=10
        )
        rows = []
        for line in result.stdout.strip().splitlines():
            vals = [v.strip() for v in line.split(",")]
            if len(vals) == len(FIELDS):
                rows.append(dict(zip(FIELDS, vals)))
        return rows
    except Exception as e:
        print(f"nvidia-smi error: {e}", file=sys.stderr)
        return []

with open("gpu_log.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["timestamp"] + FIELDS)
    writer.writeheader()
    while True:
        ts = datetime.datetime.utcnow().isoformat()
        for row in sample_gpus():
            writer.writerow({"timestamp": ts, **row})
        f.flush()
        time.sleep(5)

Expect these follow-ups

How would you replace this with pynvml to avoid the overhead of forking a subprocess every 5 seconds?

nvidia-smimonitoringgpuscriptingpython

As asked

Write a function that takes a model config (num_layers, num_heads, head_dim, dtype_bytes) and serving config (max_seq_len, max_batch_size) and returns the peak KV cache memory in bytes. Then compute it for Llama-2-70B in fp16 serving 32 concurrent requests at 4096 tokens each.

Sample answer outline

KV cache = 2 * num_layers * num_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size. For Llama-2-70B: 80 layers, 64 heads (GQA uses 8 KV heads), head_dim=128, fp16=2 bytes, 4096 tokens, 32 requests. With GQA the answer must use num_kv_heads (8) not num_heads. This is 2 * 80 * 8 * 128 * 2 * 4096 * 32 = approximately 42 GB. A strong candidate mentions GQA/MQA and its memory savings.

Reference implementation (python)

Python

def kv_cache_bytes(
    num_layers: int,
    num_kv_heads: int,  # use KV heads for GQA/MQA
    head_dim: int,
    dtype_bytes: int,
    max_seq_len: int,
    max_batch_size: int,
) -> int:
    # Factor of 2 for K and V
    return 2 * num_layers * num_kv_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size

# Llama-2-70B with GQA (8 KV heads), fp16, 4096 tokens, 32 concurrent
bytes = kv_cache_bytes(
    num_layers=80,
    num_kv_heads=8,   # GQA reduces from 64 to 8
    head_dim=128,
    dtype_bytes=2,    # fp16
    max_seq_len=4096,
    max_batch_size=32,
)
print(f"KV cache: {bytes / 1e9:.1f} GB")

Expect these follow-ups

How does grouped-query attention (GQA) reduce KV cache size, and what is the tradeoff in model quality?

kv-cachememoryinferencellamagqa

As asked

You have N=4 ranks. Ranks 0, 1, and 3 call dist.all_reduce(), but rank 2 mistakenly calls dist.broadcast() with the same tensor. Write a minimal reproducer and describe what you observe and how you fix it.

Sample answer outline

The collective type mismatch causes ranks to send incompatible protocol headers to each other; NCCL will either hang indefinitely or raise an internal error after a timeout. The reproducer should launch 4 processes with torch.multiprocessing or torchrun, have rank 2 call a different collective, and demonstrate the hang. Fix by ensuring all ranks call the same collective with the same arguments. Defensive pattern: use TORCH_DISTRIBUTED_DEBUG=DETAIL to log collective call sites.

Reference implementation (python)

Python

import torch
import torch.distributed as dist
import os

def worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    tensor = torch.ones(4).cuda(rank)
    
    if rank == 2:
        # BUG: rank 2 calls broadcast instead of all_reduce
        dist.broadcast(tensor, src=0)  # hangs or errors
    else:
        dist.all_reduce(tensor)  # all others call all_reduce
    
    print(f"Rank {rank} done: {tensor}")

# Fix: all ranks must call the same collective
# dist.all_reduce(tensor) on every rank

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    torch.multiprocessing.spawn(worker, args=(4,), nprocs=4)

Expect these follow-ups

How would you add an assertion in a multi-rank training loop to verify all ranks are about to call the same collective type?

distributednccldebuggingdeadlockpytorch

As asked

When using gradient accumulation with PyTorch DDP and accumulating over K micro-batches before each optimizer step, what is the incorrect pattern that causes unnecessary allreduce calls on every micro-batch, and how do you fix it?

Sample answer outline

DDP triggers allreduce on every backward() call by default. Without no_sync(), every micro-batch backward triggers an allreduce, wasting communication for intermediate accumulation steps. Fix: wrap all micro-batch backward passes except the last in model.no_sync() context manager. The last backward triggers the real allreduce. Also need to scale gradients by 1/K or loss by 1/K before accumulation to match full-batch semantics.

Reference implementation (python)

Python

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# WRONG: triggers allreduce on every micro-batch
def train_wrong(model, batches, optimizer):
    for i, batch in enumerate(batches):
        loss = model(batch) / len(batches)
        loss.backward()  # allreduce fires here every time!
    optimizer.step()
    optimizer.zero_grad()

# CORRECT: suppress allreduce on intermediate micro-batches
def train_correct(model: DDP, batches, optimizer):
    for i, batch in enumerate(batches):
        is_last = (i == len(batches) - 1)
        ctx = torch.distributed.algorithms.ddp_comm_hooks.default_hooks
        # Use no_sync() to suppress gradient sync on intermediate steps
        if not is_last:
            with model.no_sync():
                loss = model(batch) / len(batches)
                loss.backward()
        else:
            loss = model(batch) / len(batches)
            loss.backward()  # allreduce fires here only
    optimizer.step()
    optimizer.zero_grad()

Expect these follow-ups

Does using FSDP require a different pattern for gradient accumulation compared to DDP?

gradient-accumulationddppytorchperformancecorrectness

As asked

Write a simulation that models how increasing batch size affects per-token latency and throughput for an LLM decode step, assuming decode is memory-bandwidth-bound. Show how you would determine the optimal batch size given a latency SLA of 50 ms per token.

Sample answer outline

In the memory-bandwidth-bound regime, decode time per step scales linearly with batch size (each sequence reads the full model weights once). Throughput = batch_size / latency_per_step, which is flat in the bandwidth-bound regime since latency also scales linearly. The model should show that throughput saturates once compute-bound, and the optimal batch size for a 50 ms SLA is bandwidth / (model_bytes / 0.05). A strong answer distinguishes prefill (compute-bound) from decode (memory-bound).

Reference implementation (python)

Python

def simulate_decode_throughput(
    model_bytes: float,         # total model parameter bytes
    hbm_bandwidth: float,       # bytes/sec (e.g. 2e12 for A100)
    max_batch: int = 64,
    latency_sla_ms: float = 50,
):
    results = []
    for bs in range(1, max_batch + 1):
        # Each decode step reads model weights once per batch
        # Time scales linearly with batch size in BW-bound regime
        latency_ms = (model_bytes / hbm_bandwidth) * 1000 * bs
        throughput_tok_per_sec = bs / (latency_ms / 1000)
        meets_sla = latency_ms <= latency_sla_ms
        results.append((bs, latency_ms, throughput_tok_per_sec, meets_sla))
    return results

# 7B model in fp16 = 14 GB, A100 HBM = 2 TB/s
for bs, lat, tput, sla in simulate_decode_throughput(14e9, 2e12, latency_sla_ms=50)[:10]:
    print(f"bs={bs:3d}  lat={lat:.1f}ms  tput={tput:.0f} tok/s  sla={'OK' if sla else 'MISS'}")

Expect these follow-ups

How does tensor parallelism across multiple GPUs change the bandwidth-bound vs. compute-bound boundary?

inferencethroughputlatencybatchingmemory-bandwidth

As asked

Implement a simple ring-allreduce for a list of tensors across N ranks. Each rank should send its chunk to the next rank and receive from the previous rank, accumulating partial reductions. Show the reduce-scatter phase only.

Sample answer outline

The implementation should split the tensor into N chunks, iterate N-1 times sending the current chunk to rank+1 and receiving from rank-1, accumulating the received chunk into the local buffer. Key correctness points: non-blocking sends/receives to avoid deadlock, correct wraparound for rank 0 and rank N-1, and proper chunk indexing. A strong candidate notes that a real implementation uses CUDA-aware MPI or NCCL and avoids this Python approach in production.

Reference implementation (python)

Python

import torch
import torch.distributed as dist

def ring_reduce_scatter(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
    """Reduce-scatter phase of ring-allreduce."""
    chunks = tensor.chunk(world_size)
    # Each rank owns one chunk after reduce-scatter
    buf = chunks[rank].clone()
    
    send_rank = (rank + 1) % world_size
    recv_rank = (rank - 1) % world_size
    
    for step in range(world_size - 1):
        send_chunk_idx = (rank - step) % world_size
        recv_chunk_idx = (rank - step - 1) % world_size
        
        recv_buf = torch.empty_like(chunks[recv_chunk_idx])
        send_req = dist.isend(chunks[send_chunk_idx], dst=send_rank)
        recv_req = dist.irecv(recv_buf, src=recv_rank)
        send_req.wait(); recv_req.wait()
        
        # Accumulate into the chunk we will own
        if recv_chunk_idx == rank:
            buf += recv_buf
    
    return buf

Expect these follow-ups

How would you extend this to the all-gather phase to complete a full allreduce?

allreducedistributedringcollective-opsmpi

As asked

Write a function that takes pipeline depth P and number of micro-batches M and computes the bubble fraction (fraction of time the pipeline is not fully occupied) using the 1F1B (one forward one backward) schedule.

Sample answer outline

Bubble fraction = (P - 1) / M for the 1F1B schedule. The function should verify M >= P for a valid schedule, compute the formula, and optionally visualize the schedule as a grid. A strong candidate knows this comes from the Megatron paper and explains why increasing M reduces the bubble at the cost of more micro-batch overhead.

Reference implementation (python)

Python

def pipeline_bubble_fraction(pipeline_depth: int, num_microbatches: int) -> float:
    """
    Compute the pipeline bubble fraction for the 1F1B schedule.
    
    Args:
        pipeline_depth: number of pipeline stages P
        num_microbatches: number of micro-batches M (must be >= P)
    
    Returns:
        bubble_fraction in [0, 1)
    """
    P, M = pipeline_depth, num_microbatches
    if M < P:
        raise ValueError(f"M ({M}) must be >= P ({P}) for 1F1B to be valid")
    # Warm-up phase has (P-1) idle slots; total slots = M * P
    # Bubble slots = P - 1 (before pipeline is full) + P - 1 (drain)
    bubble_fraction = (P - 1) / M
    return bubble_fraction

# Example: 8 stages, 16 micro-batches
print(pipeline_bubble_fraction(8, 16))  # 0.4375
print(pipeline_bubble_fraction(8, 64))  # 0.109375

Expect these follow-ups

How does the interleaved 1F1B schedule (virtual pipeline stages) improve the bubble fraction?

pipeline-parallelism1f1bmegatronscheduledistributed-training

As asked

Implement a simplified forward pass of LayerNorm in OpenAI Triton that fuses the mean, variance, and normalization into a single kernel. The input is a 2D tensor of shape (N, D) and normalization is over the D dimension.

Sample answer outline

The Triton kernel should use tl.program_id to identify the row, load D elements in blocks, compute mean via tl.sum/D, compute variance similarly, then normalize in one pass, writing output. Key points: use tl.constexpr for BLOCK_D, handle non-power-of-two D with masking, and accumulate in fp32 for numerical stability before casting back to the input dtype. A strong answer discusses when this is faster than torch.nn.LayerNorm (on small D where the HBM round trip costs dominate).

Reference implementation (python)

Python

import triton
import triton.language as tl
import torch

@triton.jit
def layernorm_fwd_kernel(
    X, Y, Mean, Rstd,
    stride, N, eps,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    X_ptr = X + row * stride
    Y_ptr = Y + row * stride
    cols = tl.arange(0, BLOCK_D)
    mask = cols < N
    x = tl.load(X_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    mean = tl.sum(x, axis=0) / N
    var = tl.sum((x - mean) * (x - mean), axis=0) / N
    rstd = 1.0 / tl.sqrt(var + eps)
    y = (x - mean) * rstd
    tl.store(Y_ptr + cols, y.to(tl.float16), mask=mask)
    tl.store(Mean + row, mean)
    tl.store(Rstd + row, rstd)

Expect these follow-ups

How would you add the weight and bias (gamma, beta) parameters to this kernel without adding extra HBM reads?

tritonlayernormkernel-fusioncudaperformance

Practise these patterns on AlgoExpert

Recommended

200+ video-explained coding interview questions organised by the patterns covered on this page, with timed practice and solution walkthroughs.

Start practising

An external resource we recommend. AlgoExpert is not affiliated with us and we earn nothing from this link.

Tools to sharpen your prep

All tools

# Estimate allreduce traffic for DDP vs FSDP num_params = 7e9 bytes_per_param_bf16 = 2 # bf16 world_size = 64 # DDP: ring-allreduce sends 2*(N-1)/N * data per GPU ddp_bytes_per_gpu = 2 * (world_size - 1) / world_size * num_params * bytes_per_param_bf16 print(f"DDP allreduce per GPU: {ddp_bytes_per_gpu / 1e9:.2f} GB") # FSDP FULL_SHARD: reduce-scatter + all-gather, each ~data/N per rank per step # Each rank stores 1/N params, but communicates full gradient volume in sharded fashion fsdp_reduce_scatter = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size fsdp_all_gather = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size print(f"FSDP reduce-scatter per GPU: {fsdp_reduce_scatter / 1e9:.2f} GB")

# Roofline model calculation peak_tflops = 312e12 # A100 SXM bf16 peak_bandwidth = 2e12 # bytes/s, A100 HBM3 kernel_flops = 512e9 kernel_bytes = 8e9 arithmetic_intensity = kernel_flops / kernel_bytes # FLOPS/byte ridge_point = peak_tflops / peak_bandwidth if arithmetic_intensity < ridge_point: bound = "memory-bandwidth" roofline_time = kernel_bytes / peak_bandwidth else: bound = "compute" roofline_time = kernel_flops / peak_tflops print(f"Arithmetic intensity: {arithmetic_intensity:.1f} FLOPS/byte") print(f"Ridge point: {ridge_point:.1f} FLOPS/byte") print(f"Bound: {bound}") print(f"Roofline predicted time: {roofline_time*1000:.2f} ms")

import subprocess, csv, time, datetime, sys FIELDS = ["index", "utilization.gpu", "memory.used", "memory.total"] def sample_gpus(): try: result = subprocess.run( ["nvidia-smi", f"--query-gpu={','.join(FIELDS)}", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=10 ) rows = [] for line in result.stdout.strip().splitlines(): vals = [v.strip() for v in line.split(",")] if len(vals) == len(FIELDS): rows.append(dict(zip(FIELDS, vals))) return rows except Exception as e: print(f"nvidia-smi error: {e}", file=sys.stderr) return [] with open("gpu_log.csv", "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=["timestamp"] + FIELDS) writer.writeheader() while True: ts = datetime.datetime.utcnow().isoformat() for row in sample_gpus(): writer.writerow({"timestamp": ts, **row}) f.flush() time.sleep(5)

def kv_cache_bytes( num_layers: int, num_kv_heads: int, # use KV heads for GQA/MQA head_dim: int, dtype_bytes: int, max_seq_len: int, max_batch_size: int, ) -> int: # Factor of 2 for K and V return 2 * num_layers * num_kv_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size # Llama-2-70B with GQA (8 KV heads), fp16, 4096 tokens, 32 concurrent bytes = kv_cache_bytes( num_layers=80, num_kv_heads=8, # GQA reduces from 64 to 8 head_dim=128, dtype_bytes=2, # fp16 max_seq_len=4096, max_batch_size=32, ) print(f"KV cache: {bytes / 1e9:.1f} GB")

import torch import torch.distributed as dist import os def worker(rank, world_size): dist.init_process_group("nccl", rank=rank, world_size=world_size) tensor = torch.ones(4).cuda(rank) if rank == 2: # BUG: rank 2 calls broadcast instead of all_reduce dist.broadcast(tensor, src=0) # hangs or errors else: dist.all_reduce(tensor) # all others call all_reduce print(f"Rank {rank} done: {tensor}") # Fix: all ranks must call the same collective # dist.all_reduce(tensor) on every rank if __name__ == "__main__": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "12355" torch.multiprocessing.spawn(worker, args=(4,), nprocs=4)

import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # WRONG: triggers allreduce on every micro-batch def train_wrong(model, batches, optimizer): for i, batch in enumerate(batches): loss = model(batch) / len(batches) loss.backward() # allreduce fires here every time! optimizer.step() optimizer.zero_grad() # CORRECT: suppress allreduce on intermediate micro-batches def train_correct(model: DDP, batches, optimizer): for i, batch in enumerate(batches): is_last = (i == len(batches) - 1) ctx = torch.distributed.algorithms.ddp_comm_hooks.default_hooks # Use no_sync() to suppress gradient sync on intermediate steps if not is_last: with model.no_sync(): loss = model(batch) / len(batches) loss.backward() else: loss = model(batch) / len(batches) loss.backward() # allreduce fires here only optimizer.step() optimizer.zero_grad()

def simulate_decode_throughput( model_bytes: float, # total model parameter bytes hbm_bandwidth: float, # bytes/sec (e.g. 2e12 for A100) max_batch: int = 64, latency_sla_ms: float = 50, ): results = [] for bs in range(1, max_batch + 1): # Each decode step reads model weights once per batch # Time scales linearly with batch size in BW-bound regime latency_ms = (model_bytes / hbm_bandwidth) * 1000 * bs throughput_tok_per_sec = bs / (latency_ms / 1000) meets_sla = latency_ms <= latency_sla_ms results.append((bs, latency_ms, throughput_tok_per_sec, meets_sla)) return results # 7B model in fp16 = 14 GB, A100 HBM = 2 TB/s for bs, lat, tput, sla in simulate_decode_throughput(14e9, 2e12, latency_sla_ms=50)[:10]: print(f"bs={bs:3d} lat={lat:.1f}ms tput={tput:.0f} tok/s sla={'OK' if sla else 'MISS'}")

import torch import torch.distributed as dist def ring_reduce_scatter(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor: """Reduce-scatter phase of ring-allreduce.""" chunks = tensor.chunk(world_size) # Each rank owns one chunk after reduce-scatter buf = chunks[rank].clone() send_rank = (rank + 1) % world_size recv_rank = (rank - 1) % world_size for step in range(world_size - 1): send_chunk_idx = (rank - step) % world_size recv_chunk_idx = (rank - step - 1) % world_size recv_buf = torch.empty_like(chunks[recv_chunk_idx]) send_req = dist.isend(chunks[send_chunk_idx], dst=send_rank) recv_req = dist.irecv(recv_buf, src=recv_rank) send_req.wait(); recv_req.wait() # Accumulate into the chunk we will own if recv_chunk_idx == rank: buf += recv_buf return buf

def pipeline_bubble_fraction(pipeline_depth: int, num_microbatches: int) -> float: """ Compute the pipeline bubble fraction for the 1F1B schedule. Args: pipeline_depth: number of pipeline stages P num_microbatches: number of micro-batches M (must be >= P) Returns: bubble_fraction in [0, 1) """ P, M = pipeline_depth, num_microbatches if M < P: raise ValueError(f"M ({M}) must be >= P ({P}) for 1F1B to be valid") # Warm-up phase has (P-1) idle slots; total slots = M * P # Bubble slots = P - 1 (before pipeline is full) + P - 1 (drain) bubble_fraction = (P - 1) / M return bubble_fraction # Example: 8 stages, 16 micro-batches print(pipeline_bubble_fraction(8, 16)) # 0.4375 print(pipeline_bubble_fraction(8, 64)) # 0.109375

import triton import triton.language as tl import torch @triton.jit def layernorm_fwd_kernel( X, Y, Mean, Rstd, stride, N, eps, BLOCK_D: tl.constexpr, ): row = tl.program_id(0) X_ptr = X + row * stride Y_ptr = Y + row * stride cols = tl.arange(0, BLOCK_D) mask = cols < N x = tl.load(X_ptr + cols, mask=mask, other=0.0).to(tl.float32) mean = tl.sum(x, axis=0) / N var = tl.sum((x - mean) * (x - mean), axis=0) / N rstd = 1.0 / tl.sqrt(var + eps) y = (x - mean) * rstd tl.store(Y_ptr + cols, y.to(tl.float16), mask=mask) tl.store(Mean + row, mean) tl.store(Rstd + row, rstd)

Questions

Allreduce communication volume calculationCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Roofline model for CUDA kernel analysisCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Parse nvidia-smi output to track per-GPU utilization over timeCodingeasyCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Calculate KV cache size for a given model and batch configurationCodingeasyCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Debug a collective type mismatch causing a distributed deadlockCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Gradient accumulation with DDP: correct vs. incorrect patternsCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Throughput vs. latency tradeoff in LLM serving batchingCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Implement a basic ring-allreduceCodinghardOccasional

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Simulate the 1F1B pipeline schedule bubble fractionCodingmediumOccasional

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Write a fused LayerNorm kernel in TritonCodinghardOccasional

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Related questions

Roofline model for CUDA kernel analysis

Parse nvidia-smi output to track per-GPU utilization over time

Calculate KV cache size for a given model and batch configuration

Debug a collective type mismatch causing a distributed deadlock

More ai infrastructure engineer topics

Tools to sharpen your prep

Questions

Allreduce communication volume calculationCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Roofline model for CUDA kernel analysisCodingmediumCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Parse nvidia-smi output to track per-GPU utilization over timeCodingeasyCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Calculate KV cache size for a given model and batch configurationCodingeasyCommon

As asked

Sample answer outline

Reference implementation (python)

Expect these follow-ups

Debug a collective type mismatch causing a distributed deadlockCodingmediumCommon