ML coding interview practice

Write attention, KV cache, schedulers and loss functions from scratch, checked by tests in your browser. Free to start: no credit card, no install. Pro is $19.99 a month or $129 a year.

• GQA Attention · Hard: Implement grouped-query attention with a bottom-right causal mask.
• KVCache · Medium: Implement append and read for a preallocated KV cache.
• Top-k / Top-p Filter · Medium: Filter logits with top-k followed by top-p.
• Slot Mapping · Easy: Map logical token positions to slots in non-contiguous physical KV blocks.
• Store into KV Cache · Easy: Write K/V in place from a slot mapping and skip padding slots.
• Paged Attention Decode · Hard: Read a valid context from shuffled physical blocks and run one GQA decode step.
• Variable-Length Attention · Hard: Implement varlen attention with cu_seqlens and a bottom-right causal mask.
• Prepare Prefill Metadata · Medium: Build input ids, positions, cu_seqlens, slot mappings and block tables for prefill.
• Prepare Decode Metadata · Easy: Build positions, slots and context lengths for one decode token per sequence.
• BlockManager with Prefix Caching · Hard: Implement block allocation, reference counts, chained prefix hashes and reclamation.
• Assemble Attention.forward · Hard: Combine KV writes, varlen prefill and paged decode into a complete attention forward.
• Simulate Static Batching · Easy: Simulate fixed batches where short requests wait for the longest request.
• Simulate Continuous Batching · Medium: Admit new requests after each iteration and track continuous batching execution.
• Implement the nano-vllm Scheduler · Hard: Implement prefill, decode, preemption and postprocess under tight KV capacity.
• Token Budget Mixed Scheduling · Medium: Schedule decode first, then chunk prefills under one token budget.
• Chunked Prefill Attention · Medium: Process a prompt in chunks while matching one-shot causal attention.
• RadixTree Prefix Matching · Hard: Implement edge splitting, insertion and longest-prefix matching in a compressed radix tree.
• Prefix Cache Locking and LRU Eviction · Medium: Propagate lock references and safely evict unlocked KV leaves by LRU.
• Stable Softmax and LogSumExp · Medium: Implement stable softmax and logsumexp, including extreme logits and all -inf rows.
• Cross-Entropy with Ignore Index · Medium: Compute mean cross-entropy and its logits gradient while skipping ignored labels.
• LayerNorm and RMSNorm · Hard: Implement forward and backward passes for two Transformer normalization layers.
• AdamW Update Step · Medium: Implement one AdamW update with bias correction and decoupled weight decay.
• Gradient Accumulation and Global Clipping · Medium: Combine microbatch gradients by valid-token count, then clip their global norm.
• Warmup + Cosine Learning Rate · Easy: Implement a learning-rate schedule with linear warmup followed by cosine decay.
• Multi-Head Causal Attention · Medium: Split Q/K/V into heads, apply scaled causal attention, and merge the heads.
• Pre-LN Transformer Block · Hard: Combine LayerNorm, causal multi-head attention, a GELU MLP, and two residual paths.
• GPU Peak FLOPS · Easy: Compute peak TFLOPS from SMs, units, FMAs per clock and clock speed, and convert datasheet sparse peaks to dense.
• Roofline Bound · Medium: Compute arithmetic intensity and the attainable ceiling, and decide whether a kernel is compute- or bandwidth-bound.
• Ring All-Reduce Time · Medium: Separate bidirectional from one-way bandwidth and estimate ideal gradient-sync time across GPU counts and links.