Meta PyTorch interview problems — TorchLeet

Implement KNN in PyTorch

Build k-nearest neighbors classification using vectorized distance computation and top-k selection on tensors.

Easy

Implement Logistic Regression with Gradient Descent

Build binary logistic regression from scratch with sigmoid, binary cross-entropy, and manual gradient descent updates.

Easy

Implement Softmax from Scratch

Build numerically stable softmax using the log-sum-exp trick, handling overflow and underflow in raw tensor math.

Easy

Implement Beam Search for LLM Decoding

Maintain and expand top-scoring partial sequences at each decoding step with length normalization and early stopping.

Medium

Implement Knowledge Distillation

Train a smaller student model to match a larger teacher's soft predictions using temperature-scaled KL divergence loss.

Medium

Implement KV Cache for Autoregressive Generation

Cache key and value tensors from previous timesteps so each new token only computes attention over one new position.

Medium

Implement LoRA on a Linear Layer

Inject trainable low-rank decomposition matrices (A, B) into a frozen linear layer for parameter-efficient fine-tuning.

Medium

Implement a Swin Transformer Block (Shifted Window Attention)

Build a Swin Transformer block from scratch: window partition, shifted-window attention masks, and a learned relative position bias.

Hard

Implement Continuous Batching for LLM Inference

Dynamically slot sequences in and out of a running batch as they finish, maximizing throughput without padding waste.

Hard

Implement DPO Loss from Scratch

Compute the Direct Preference Optimization loss that trains a policy directly from preference pairs without a reward model.

Hard

Implement Gradient Checkpointing

Trade compute for memory by recomputing intermediate activations during backward instead of storing them all.

Hard

Implement PPO for RLHF

Build Proximal Policy Optimization with clipped surrogate objective, value function baseline, and KL penalty for RLHF.

Hard

Implement Vision Transformer + MAE Pretraining

Build ViT with masked autoencoder pretraining — randomly mask patches, encode visible ones, decode to reconstruct.

Hard

Implement FlashAttention-2 in Triton

Write the tiled, fused attention kernel that computes exact attention in O(n) memory using online softmax in Triton.

Expert

Implement FSDP from Scratch

Build Fully Sharded Data Parallel: shard parameters across GPUs, all-gather before forward, reduce-scatter gradients after backward.

Expert

Implement Ring Attention for Long Contexts

Distribute attention computation across GPUs in a ring topology, enabling context lengths that exceed single-GPU memory.

Expert

Write a Fused Softmax Kernel in Triton

Write a GPU kernel in Triton that fuses the softmax computation into a single pass, eliminating intermediate memory reads.

Expert