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. ```hint Test what a causal model guarantees Some attention bugs leave every shape unchanged. Think about which input positions are allowed to influence each output position of a causal language model, and turn that property into a check. ``` #### 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. ```hint Track absolute positions When only the newest token passes through the model, consider which cached positions its query may attend to and which position index the token should receive. ``` #### 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. ### What a Strong Answer Covers - Precise reasoning about scaled dot-product attention and its tensor shapes, applied consistently to both debugging and caching. - Verification through small deterministic tests rather than by inspecting generated text. - Correct inference-time settings during generation, such as disabled dropout and no gradient tracking. - Clear trade-offs around cache memory growth, buffer allocation, and context-length limits. ### Follow-up Questions 1. How would you support batched generation for prompts of different lengths with a KV cache? 2. How does the cache change with rotary position embeddings, or with grouped-query attention in which several query heads share one key/value head? 3. How would you avoid reallocating the cache at every step, and what should happen when generation reaches the maximum context length? 4. How would you reorder the cache during beam search?