• /home
  • /about
  • /posts
  • /resume

Reference Videos: Here These are rough short notes, attuned to my needs w.r.t to the topic.

Index

Some Terms

Tokenization

Every token maps to an embedding (trainable):

Positional Embeddings & RoPE

The Math Behind RoPE: The Relative Distance Function

When we talked about Rotary Position Embeddings (RoPE) turning absolute positions into relative distances, the “Function” in that relationship relies on the geometry of complex numbers.

To see how the absolute positions (let’s call them $m$ for the query token and $n$ for the key token) collapse into a relative distance $(m-n)$, we can map the token vectors onto a 2D complex plane.

Instead of adding a position vector, RoPE multiplies the query and key vectors by a complex rotation. If $\theta$ is the base angle frequency, we rotate the query vector $q$ by angle $m\theta$ and the key vector $k$ by angle $n\theta$:

\[q_m = q e^{im\theta}\] \[k_n = k e^{in\theta}\]

The attention score is calculated by taking the dot product (the inner product) of these two vectors. In the complex plane, this looks like:

\[\langle q_m, k_n \rangle = \text{Re}(q_m k_n^*)\]

Substitute our rotated vectors into the equation:

\[\langle q_m, k_n \rangle = \text{Re}(q e^{im\theta} k^* e^{-in\theta})\] \[\langle q_m, k_n \rangle = \text{Re}(q k^* e^{i(m-n)\theta})\]

That right there is the magic. The absolute positions $m$ and $n$ vanish as independent variables. The attention score is now strictly a function of the original word vectors ($q$ and $k$) and the relative distance between them $(m-n)$. This elegant mathematical property is exactly what allows modern models to extrapolate beyond their trained context windows.

The Elephant in the Room - Attention

When text enters the model, each word is initially assigned a high-dimensional vector - an embedding, where directions in space correspond to semantic meaning. However, these initial embeddings are just a basic lookup table — the word “mole” gets the exact same vector whether the text says “true mole,” “mole of carbon dioxide,” or “biopsy of the mole.”

The attention mechanism exists to fix this. It allows the embeddings to pass information to each other, updating their vectors so they bake in the surrounding context.

The Three Pillars: Queries, Keys, and Values

To figure out how words should update one another, a single “head” of attention relies on three matrices full of tunable parameters that act on the embeddings.

Think of an example phrase like “a fluffy blue creature”, where adjectives need to update the meaning of a noun:

Key insight: The attention mechanism calculates a complete grid of every possible Key-Query combination, producing a massive matrix of relevance scores between every single word in the sequence.

Refining the Attention Pattern

Before those relevance scores can be used to update the embeddings, two critical mathematical steps happen to the grid:

Because this grid maps every word against every other word, its size scales quadratically — $O(N^2)$ with the context size. This is why larger context windows are historically so computationally expensive.

Multi-Headed Attention & Scale

Everything above describes a single head of attention, which might learn just one contextual rule (like adjectives updating nouns). To capture the true complexity of language, transformers use multi-headed attention:

Training vs. Inference

How Training Works in Practice

Because the self-attention mechanism processes the entire context window simultaneously, a single forward pass over a sequence (e.g., 2,048 tokens) yields 2,048 distinct output vectors at the final layer. Instead of just calculating the loss on the very last token, the training pipeline calculates the cross-entropy loss for every position at the same time.

If the training sequence is ["The", "cat", "sat", "on", "the", "mat"], the model outputs a vector for each word simultaneously.

Why this matters: This is the core reason why Transformers are so efficient to train on GPUs compared to older architectures. By using every single column to predict the token that comes immediately after it, a single batched operation provides 2,048 different training signals rather than just one.

How Inference Works in Practice

The GPU allocates a matrix sized to the model’s maximum context window, filling unused space with padding tokens. An attention mask ensures the model only calculates attention for the active tokens.

During generation, we only look at the output vector of the last active token in the sequence (the current “final column”) to predict the new unknown token. That final vector contains the context of the entire sequence and is passed through the unembedding matrix to predict the $N+1$ token. Once predicted, that new token is appended to the input, and the process repeats.

The <EOS> Token: The autoregressive loop does not run indefinitely. During training, documents are terminated with an End of Sequence (<EOS>) token. When the final column outputs <EOS> with the highest probability, the software wrapper breaks the loop and finalizes the output.

Context Windows & Output Limits

Input vs. Output Caps: The KV Cache Dynamic

From a pure storage perspective, 1 token = 1 token. An input token and an output token occupy the exact same amount of VRAM in the KV cache.

However, output limits are capped much more heavily due to system dynamics:

KV Caching

Re-multiplying the entire matrix for historical tokens every time a new word is added wastes massive compute. The system caches the intermediate mathematical representations (Keys and Values) of past tokens in GPU memory (VRAM). For a new token, the GPU only calculates its specific values and pulls the rest from the cache.

To walk through a concrete example: if you feed the model a 3-token input ["The", "cat", "sat"], the model processes 3 columns simultaneously:

If Column 3 predicts "on", you take "on", append it to the input array so it becomes ["The", "cat", "sat", "on"], and send the whole sequence back into the model to predict the 5th token.

Because the Keys and Values of earlier tokens like "The", "cat", and "sat" do not change when "on" is added to the end of the sequence, the system caches those intermediate calculations in GPU memory. The server doesn’t actually recalculate the entire sequence — it only runs the matrix multiplications for the newly added token and pulls the historical context directly from the memory cache.

The VRAM Wall: Storing the KV cache for large context windows (or thousands of concurrent users) consumes massive amounts of VRAM, making memory capacity and memory bandwidth the primary hardware bottlenecks. Managing this cache efficiently across thousands of concurrent users is currently one of the hardest infrastructure challenges in deploying high-throughput AI systems.

FlashAttention: Defeating the $O(N^2)$ VRAM Monster

VRAM and memory bandwidth are the true physical bottlenecks of the Transformer. FlashAttention is currently the most important engineering breakthrough in the industry because it directly attacks how the GPU routes that memory.

To understand FlashAttention, you have to look at the physical architecture of a GPU, which has a strict memory hierarchy:

  1. SRAM (Compute Cache): Incredibly fast, located right next to the compute cores, but incredibly tiny (e.g., 20MB per GPU).

  2. HBM (VRAM): The main memory pool (e.g., 80GB on an H100). It is huge, but reading/writing to it is remarkably slow compared to the cores’ calculation speed.

The Problem with Standard Attention

In standard attention, the GPU calculates the massive $N \times N$ attention matrix (Query matrix multiplied by Key matrix). Because it is so huge, it cannot fit in the SRAM. The GPU is forced to write this massive $N \times N$ matrix to the slow HBM, read it back to apply the Softmax function, write it back to HBM, and read it again to multiply it by the Value matrix.

The GPU cores spend most of their time sitting idle, waiting for the massive matrix to travel back and forth from the HBM. This is known as the “Memory Wall.”

The FlashAttention Solution

Needs further reading- https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad

To be added..