Explain Transformer Internals and Implement Top-p Decoding

Context

In autoregressive language models, text generation proceeds token-by-token. As an ML Engineer optimizing inference pipelines, you must thoroughly understand both the internal mechanics of decoder-only transformers and the stochastic decoding strategies used to sample next tokens.

Part 1: Conceptual Architecture

Be prepared to address the following system and algorithmic mechanics:

1. KV Caching: Why is KV caching essential during autoregressive generation? What is the asymptotic memory complexity per request with context length L, hidden dimension d_model, number of layers n_layers, and batch size B? How does Multi-Query Attention (MQA) or Grouped-Query Attention (GQA) alter this footprint compared to standard Multi-Head Attention (MHA)?
2. Rotary Position Embeddings (RoPE): How does RoPE incorporate relative positional information into query and key representations without adding explicit absolute position embeddings to input tokens?

Part 2: Implementation Task

Implement a function sample_top_p that takes raw next-token logits from a transformer decoder and samples an index according to nucleus (top-p) filtering with temperature scaling.

Requirements

Given:

• logits: A 1D NumPy array or PyTorch tensor of shape (vocab_size,) containing unnormalized scores for the next token.
• temperature: A float > 0.0 controlling distribution entropy. Lower values make the distribution peakier; higher values flatten it.
• top_p: A float in (0.0, 1.0] representing the cumulative probability threshold for nucleus sampling.

Your function should:

1. Scale the logits by temperature.
2. Convert logits to probabilities using a numerically stable softmax.
3. Sort probabilities in descending order and identify the smallest set of tokens whose cumulative probability exceeds or equals top_p (ensuring at least one token is kept).
4. Zero out (or remove) all tokens outside this nucleus set.
5. Renormalize the remaining probabilities so they sum to 1.0.
6. Sample a single token index according to this filtered distribution.

Function Signature

import numpy as np

def sample_top_p(logits: np.ndarray, temperature: float = 1.0, top_p: float = 0.9) -> int:
    pass

Example

logits = np.array([2.0, 1.0, 0.1, -1.0])
temperature = 1.0
top_p = 0.75

# Softmax of logits ~ [0.584, 0.215, 0.087, 0.029] (normalized)
# Sorted probabilities: 0.584 (idx 0), 0.215 (idx 1), 0.087 (idx 2), 0.029 (idx 3)
# Cumulative probabilities: [0.584, 0.799, 0.886, 0.915]
# Nucleus cutoff includes idx 0 and idx 1 (cumulative sum 0.799 >= 0.75)
# Renormalized probabilities for [idx 0, idx 1]: [0.584 / 0.799, 0.215 / 0.799] ~ [0.731, 0.269]

token = sample_top_p(logits, temperature=1.0, top_p=0.75)
# token will be either 0 or 1