Katkat
Back

Debug Transformer and Add KV Cache

OpenAIMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteMachine Learning

Last updated

This problem tests your understanding of transformer architecture by requiring you to fix a broken attention mechanism and implement a Key-Value cache to optimize autoregressive inference.

Problem Statement

You are working on a minimalist GPT-style implementation. Your task is to identify and fix a bug in the provided self-attention block and implement a KV-cache mechanism to improve inference efficiency.

Part 1: Debugging

The following forward pass for a Multi-Head Attention block is meant to implement causal self-attention. Identify why it currently fails to enforce causality and how it incorrectly scales the attention scores.

def forward(q, k, v):
    # q, k, v shape: (batch_size, seq_len, head_dim)
    scores = torch.matmul(q, k.transpose(-2, -1)) # Error 1 here
    probs = F.softmax(scores, dim=-1)
    return torch.matmul(probs, v)

Part 2: KV-Cache Implementation

Modify the inference loop to support a KV-cache. Instead of recomputing k and v for the entire sequence at every step, your model should accept a past_key_values tuple, append only the current token's keys and values, and return the updated cache.

Constraints

  • batch_size: 1
  • head_dim: 64
  • Assume q is generated only for the current token (sequence length 1 during inference).

Expected Output

Your implementation should provide a function forward_with_cache(q, k, v, past_kv=None) that returns (output, new_kv).

Hints

Hint 1

Consider how the softmax operation behaves if the upper triangular portion of the attention matrix is not masked out with negative infinity.

Hint 2

Recall that the attention score requires scaling by the square root of the head dimension to prevent gradient saturation.

Hint 3

When using a KV-cache, the attention mechanism only needs to compare the current query against the concatenated keys stored from all previous time steps.

Follow-up questions

  • How does the memory complexity change with a KV-cache compared to standard full-sequence computation?
  • What are the implications of using Grouped Query Attention (GQA) on the size and structure of the KV-cache?
  • How would you handle cache eviction if the context window exceeds the allocated memory?