Debug a Transformer Implementation and Implement a KV Cache for Decoding

This ML coding exercise centers on a Transformer language model. You are handed an existing Transformer implementation into which four or five bugs have been deliberately planted. Find and fix the bugs, then implement key-value (KV) caching from scratch so that autoregressive generation reuses the attention keys and values already computed for earlier tokens instead of reprocessing the whole sequence at every step.

The buggy code itself is not reproduced here. Unless told otherwise, assume a standard decoder-only language model: token and position embeddings, a stack of blocks that each apply masked multi-head self-attention and a position-wise feed-forward network with residual connections and layer normalization, then a final projection to vocabulary logits. Inside attention, tensors have shape (batch, heads, sequence, head_dim).

Constraints and Clarifications

• The planted bugs are not identified for you; locating them is part of the task.
• Work in the tensor library the provided code uses, but write the cache yourself rather than relying on a library's built-in generation cache.
• Cached generation must produce the same logits, within floating-point tolerance, as running the full sequence through the model without a cache.

Clarifying Questions

• Does the code currently raise an error, or does it run and produce wrong outputs or a training loss that behaves unexpectedly?
• Is the model decoder-only, or is there also an encoder whose cross-attention keys and values could be cached?
• Which position scheme does the model use, such as learned absolute embeddings or rotary embeddings?
• Must generation support a batch of prompts with different lengths, and what is the maximum context length?

Part 1 — Find and Fix the Planted Bugs

Because the original code is not shown, walk through how you would find the planted bugs efficiently: which parts of the forward pass and the training loss you would inspect, what each likely bug looks like and how you would fix it, and which small tests would expose a bug even when every tensor shape is correct and the code runs without an error.

What This Part Should Cover

• A systematic pass through the attention computation (projections, head splitting and merging, scaling, masking, softmax), the block wiring, and the next-token loss.
• Targeted tests that reveal bugs that do not change tensor shapes.
• A concrete fix for each bug found and confirmation that the fixes work together.

Part 2 — Implement a KV Cache From Scratch

Change the model so that generation runs in two phases: process the prompt once, then generate one token at a time while each layer reuses the keys and values it computed for earlier positions. Write the cache data structure, the changes to the attention forward pass, the position handling, and the generation loop. Then explain how you would verify that cached and uncached generation agree, and analyze the compute and memory cost.

What This Part Should Cover

• Per-layer storage and update of keys and values with correct tensor shapes.
• Correct masking and position offsets when the number of queries differs from the number of keys.
• A test comparing cached generation against full recomputation.
• Time and memory complexity with and without the cache.