Review a Compact GPT-Style Transformer Implementation

Scenario

You are reviewing an engineer's compact, decoder-only autoregressive language model implementation written in PyTorch. The model will be used for both training with cross-entropy loss and real-time generation via key-value (KV) caching.

During testing, the team observed two major issues:

1. Training loss plateaus early at an unusually high perplexity.
2. In autoregressive inference, generating tokens one-by-one with a KV cache produces different probability distributions than passing the full generated prompt in a single forward pass.

Provided Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, past_kv=None):
        B, T, C = x.shape
        qkv = self.qkv_proj(x)  # (B, T, 3 * C)
        q, k, v = torch.chunk(qkv, 3, dim=-1)
        
        # Reshape to (B, n_heads, T, head_dim)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=-2)
            v = torch.cat([past_v, v], dim=-2)
        new_kv = (k, v)
        
        # Attention computation
        scores = torch.matmul(q, k.transpose(-2, -1))
        mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0).unsqueeze(0)
        scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out), new_kv

class MiniGPT(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx, past_kv=None):
        x = self.tok_emb(idx)
        x, new_kv = self.attn(x, past_kv=past_kv)
        logits = self.head(x)
        return logits, new_kv

def compute_loss(logits, targets):
    # logits: (B, T, vocab_size), targets: (B, T)
    return F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

Tasks

1. Trace Shapes & Numerical Behavior: Identify the missing scaling factor in attention scores and explain the numerical instability that occurs as head_dim grows.
2. Identify the Causal Mask & Cache Mismatch: When past_kv is provided during step-by-step decoding (T = 1, but k.shape[-2] = total_len), explain why the attention computation breaks or produces incorrect attention distributions.
3. Fix the Loss Alignment: Identify the target alignment bug in compute_loss during training and rewrite the function to correctly compute next-token prediction loss.
4. Parity Verification: Outline how you would construct a unit test to verify that autoregressive decoding with a KV cache produces outputs identical to an uncached forward pass within floating-point tolerance.