Attention is one of the two major computational sublayers in a Transformer block; the other is the feed-forward network. Normalization, residual connections, positional mechanisms, and other details sit around these two operations. Part 1 covered the attention half of that pair: how relationships between tokens get scored, how far a token can see, which direction information flows, and how Q, K, and V get split across heads.
The next step is to look at how a real model might take in information and process it, because a few things change when training a model at scale.
Attention, as defined so far, is powerful but incomplete. It has no built-in sense of where a token sits in the sequence, no mechanism for skipping re-computation as generation proceeds one token at a time, and no answer for how the underlying matrix math survives contact with real hardware once sequences get long or the input stops being text altogether. Modern transformer engineering has settled on some clever, fairly durable answers to each of those gaps.
While research is large and still active, and new papers continue to add variants to each of these problem areas, this article covers four techniques that sit among the most widely adopted answers to date: RoPE, the KV cache, Flash attention, and cross-modal attention.
Consider two short sentences built from the same three content words:
The cat ate the fish. The fish ate the cat.
The meaning is reversed, even though the two sentences use identical tokens. Whatever mechanism scores relationships between those tokens needs some signal that distinguishes “cat before ate” from “fish before ate,” or the two sentences become indistinguishable to it. Attention, as defined so far, has no such signal built in.
Here is that problem stated more precisely, hiding in plain sight in the scaled dot-product formula from Part 1:
Attention(Q, K, V) = softmax( (Q · Kᵀ) / √dₖ ) · V
Nothing in this equation cares about order. If the tokens in a sentence were scrambled and presented in a different order, the content-dependent Q·Kᵀ interactions would remain the same for the same token pairs, but their positions in the resulting attention matrix would be rearranged.
Transformers need to inject position information somehow. The original Attention Is All You Need paper did this by adding a fixed sinusoidal vector to each token’s embedding, and later models used learned position embeddings instead. Both provide positional information, but they encode position differently from RoPE: token 5 receives information associated with its absolute position, as does token 500, while relative-distance relationships have to be inferred from those positional representations.
Rotary Position Embedding (RoPE) takes a different approach: instead of adding position information, it rotates the query and key vectors by an angle that depends on position.
Take a pair of dimensions from a vector, (x₂ᵢ, x₂ᵢ₊₁), and rotate them by an angle θ that grows with the token's position m:
θ = m · base^(-2i/d)
A simple picture helps here. Think of each pair of dimensions as a clock hand. Position 0 leaves the hand untouched; each later position nudges the hand forward by a fixed angle:
Position 0 Position 1 Position 2 ↑ ↗ → | (0°) (θ) (2θ)
Each step adds another angular increment.
Different dimension-pairs rotate at different speeds, encoding positional information.
The part that makes RoPE genuinely useful is what happens at the scoring step. When the query is rotated by its position m and the key is rotated by its position n, and the two are then combined via dot product, the absolute positions m and n do not vanish from the computation; instead, their difference appears as a relative rotation of (n − m).
(Rₘ Q) · (Rₙ K) = Qᵀ Rₙ₋ₘ K
Two concepts from the formula above:
Token at position 50 Token at position 3050 ↓ ↓ rotate by 50θ rotate by 3050θ ↓ ↓ └──────────── dot product ─────────────┘ ↓positional component reflects the relative offset: 50 − 3050 = −3000
RoPE’s relative-position property does not make it immune to context-length extrapolation problems. When a model is evaluated substantially beyond the positional range represented during training, its learned use of the rotary frequencies can degrade, motivating techniques such as position interpolation, frequency rescaling, or YaRN.
import torchdef apply_rope(x, base=10000): # x: (seq_len, dim), applied as a preprocessing step to Q or K. seq_len, dim = x.shape positions = torch.arange(seq_len).unsqueeze(1).float() # (seq_len, 1), this is m freq_idx = torch.arange(0, dim, 2).float() # (dim/2,), this is 2i # inv_freq[i] = base^(-2i/d) inv_freq = 1.0 / (base ** (freq_idx / dim)) # (dim/2,) theta = positions * inv_freq.unsqueeze(0) # (seq_len, dim/2), outer product: m * base^(-2i/d) cos, sin = torch.cos(theta), torch.sin(theta) # Note: interleaving even/odd dimensions matches the original RoPE # paper. x1, x2 = x[..., 0::2], x[..., 1::2] # even/odd dims rotated = torch.empty_like(x) rotated[..., 0::2] = x1 * cos - x2 * sin rotated[..., 1::2] = x1 * sin + x2 * cos return rotated
Note that RoPE is applied to Q and K before the dot product, never to V, since V carries content, not position.
Recall causal self-attention from Part 1: each token can only attend to itself and everything before it. That constraint has a useful side effect. Once a token’s key and value vectors are computed, they never change, because nothing that comes after a token can influence its own K or V, since attention only flows backward.
Consider what this means for text generation, one token at a time. Naively, at every single step, the model would recompute Q, K, and V for the entire sequence so far, run full attention, and keep only the last token’s output. Five hundred tokens into generation, that means recalculating the K and V for all 499 earlier tokens at every single step, even though those values are identical to what was computed the step before. It is comparable to rewriting the first five hundred pages of a book from scratch every time a single new sentence gets added to page 501.
The KV cache simply avoids that redundancy. Keys and values for already-processed tokens are stored once and reused. On each new step, the model computes Q, K, and V only for the newest token, appends that K and V to the cache, and attends the new query against everything stored so far.
Step 1: "The" cache: K₁ V₁
Step 2: "The cat" cache: K₁ V₁ | K₂ V₂ ← K₁,V₁ reused, not recomputed
Step 3: "The cat sat" cache: K₁ V₁ | K₂ V₂ | K₃ V₃ ← only K₃,V₃ are new
Each step appends one token’s K/V representation.
python
class KVCache: def __init__(self): self.k_cache = None self.v_cache = None def update(self, new_k, new_v): # new_k, new_v: (1, dim) for the newest token if self.k_cache is None: self.k_cache, self.v_cache = new_k, new_v else: # Pedagogical only. Production caches preallocate # or use block-based storage. self.k_cache = torch.cat([self.k_cache, new_k], dim=0) self.v_cache = torch.cat([self.v_cache, new_v], dim=0) return self.k_cache, self.v_cachedef generate_step(q_new, k_new, v_new, cache): k_full, v_full = cache.update(k_new, v_new) d_k = q_new.size(-1) scores = (q_new @ k_full.transpose(-2, -1)) / (d_k ** 0.5) weights = torch.softmax(scores, dim=-1) return weights @ v_full
Without a cache, the attention portion of generation repeatedly recomputes the entire prefix at every decoding step. At step i, this can require O(i²) attention work, and summed across n decoding steps, the cumulative attention work reaches O(n³).
By using a KV cache during decoding, a single new query token only needs to attend against the i keys already stored in memory. That drops a single decoding step down to O(i) compute, turning the overall generation process into O(n²).
But that saving introduces a new memory-bandwidth bottleneck. Every layer and every attention head needs its own set of K and V vectors saved for every generated token. These vectors get stored in the GPU’s high-bandwidth memory (HBM, effectively its VRAM). HBM stores the large tensors, while the GPU’s compute units perform the matrix operations using data brought through the GPU’s memory hierarchy, including much smaller and faster on-chip memories such as registers and shared memory/SRAM. At each decoding step, the GPU must repeatedly read the accumulated KV cache through the memory hierarchy so the new query can compare against the stored keys and combine their values. With a small batch size, this means repeatedly moving large amounts of cached data through the GPU’s memory hierarchy for relatively little new arithmetic.
This pressure motivated multi-query and grouped-query attention. Multi-query attention (MQA) forces every head to share a single K/V pair, shrinking the cache drastically at the expense of some model capacity. Grouped-query attention (GQA) strikes a balance, sharing K/V pairs within small groups of heads instead of across all of them or none of them. A smaller cache means less data shuttling between HBM and SRAM on every single token, freeing the GPU to run larger batch sizes and deliver much higher throughput.
Part 1 noted that global attention costs O(n²) in both compute and memory: the attention matrix has one entry for every pair of tokens. For a sequence of 32,000 tokens, that matrix has over a billion entries, and a naive implementation has to materialize that entire matrix in GPU memory before it can run SoftMax on it.
GPUs can do the floating-point operations fast enough; the slow part is moving a matrix that large between two tiers of GPU memory: the large, comparatively slow high-bandwidth memory (HBM) where tensors normally live, and the small, very fast on-chip SRAM where the GPU actually performs computation. Memory bandwidth, not compute, is often what limits the whole operation.
Flash attention never materializes the full n×n matrix in HBM. It avoids writing intermediate n×n attention scores back out to HBM. Blocks of Q, K, and V are loaded through the GPU memory hierarchy into on-chip memory, where tiled matrix multiplication, scaling, and online softmax can be performed without materializing the full attention matrix in HBM. Q is processed in row blocks while K and V are processed along their sequence dimension, creating nested Q-block and K/V-block loops.
A useful mental picture: rather than printing out an enormous spreadsheet in full just to perform one calculation, the tiles are kept in fast, on-chip memory and processed a small block at a time, with the matrix multiplication and softmax steps fused together so intermediate results never have to round-trip through slower memory.
The mechanism that makes this possible is the online softmax. Ordinarily, softmax needs to see every score in a row before normalizing any of them, since the normalizing sum depends on all of them. Flash attention tracks a running maximum and sum, rescaling the output computed so far whenever a new, larger maximum appears.
The online-softmax result is mathematically equivalent to standard attention. Flash attention is exact and can recompute needed attention values during the backward pass instead of storing the full n×n matrix. That IO-conscious design, on both passes, is what “IO-aware” in the algorithm’s name refers to: it is built around the cost of moving data between HBM and SRAM, not just the number of floating-point operations it performs. This dual tiling and recomputation are also why a correct, efficient Flash attention implementation is considerably harder to get right than the single-loop sketch above suggests.
A standard implementation that materializes the full attention matrix moves roughly O(n²) worth of data between HBM and SRAM, since the whole n×n score matrix gets written out and read back in. Flash attention’s tiling reduces HBM traffic by keeping intermediate attention blocks on chip instead of repeatedly materializing them in HBM. Its IO complexity depends on sequence length, head dimension, and available SRAM; importantly, Flash attention does not turn the quadratic attention computation into linear-time computation. Its major memory advantage is that its memory footprint grows linearly with sequence length rather than requiring an O(n²) attention matrix. The FLOP count changes little; the main savings come from reducing memory traffic.
Part 1 introduced cross-attention as “Q comes from one sequence, K and V come from another,” a decoder generating text while attending to an encoder’s output. That same pattern is what lets attention step outside text entirely.
The trick is that an image can be turned into something that looks like a sequence of tokens. In patch-based vision models, an image is split into fixed-size patches that are projected into vectors. Each patch becomes a vector, and a grid of patches becomes a “sequence” of patch embeddings, no different in shape from a sequence of word embeddings.
Once represented as tokens, image features can participate in cross-attention, allowing text queries to incorporate visual information.
Text query: "cat" │ ▼ compares against every patch key │ ▼p1:0.02 p2:0.03 p6:0.71 p7:0.15 ... p16:0.01 ↑ strongest attention weight falls on a visual token
python
def cross_modal_attention(text_q, image_k, image_v): # text_q: (text_len, d_model), queries from the text decoder # image_k: (num_patches, d_model), keys from image patch embeddings # image_v: (num_patches, d_model), values from image patch embeddings d_k = text_q.size(-1) scores = (text_q @ image_k.transpose(-2, -1)) / (d_k ** 0.5) weights = torch.softmax(scores, dim=-1) return weights @ image_v, weights
Related attention mechanisms are widely used in image captioning, text-conditioned image generation, and multimodal assistants, although the direction of the Q, K, and V projections and the surrounding architecture can differ substantially between systems. The same idea extends to audio in older encoder-decoder systems such as Whisper, where waveform segments or spectrogram patches take the place of image patches as the sequence being attended over via cross-attention. Some audio-language models instead tokenize audio with a neural codec and process those tokens alongside text using self-attention.
Two simplifications are worth flagging in the walkthrough above. First, production vision-language models rarely feed raw patch embeddings straight into the language model. A dedicated vision encoder (CLIP or SigLIP, for instance) produces the patch representations first, and those typically pass through additional pooling or a small projection network that compresses the patch count substantially before anything reaches the language model, both to control compute cost and to align the vision encoder’s representation space with the language models. Second, cross-attention in the Part 1 sense, where Q comes from text and K/V come from a separate encoder, is only one way to combine modalities. A number of multimodal architectures instead project visual representations into a form that can be processed alongside language tokens, allowing self-attention to mediate interactions between the modalities. Other architectures retain dedicated cross-attention or use an intermediate module to connect the vision and language representations.
The river carved through the canyon for centuries. Eventually, the canyon began to shape how the river flowed.
A production model can combine these techniques around the same attention operation. When it reaches the second occurrence of “canyon,” self-attention connects that token back to the first one, however many tokens separate them, exactly as in Part 1. What changes in production is everything running around that connection:
They are the layers built around the core operations from Part 1: compare Q against K, scale, mask, turn scores into weights with softmax, and use those weights to combine V. RoPE makes attention position-aware; the KV cache reduces redundant computation; FlashAttention reduces memory traffic; and cross-modal attention extends the same mechanism beyond text.
One caution is worth carrying forward. A single attention weight is not a window into a model’s reasoning. It is a learned signal that shapes what gets blended into a token’s representation, not a labeled trace of a thought process. Attention can still play causal roles, but establishing those roles requires systematic circuit analysis, not a glance at one high attention weight.
Together, these ideas show why attention is more than the equation suggests. The core operation is simple; making it useful in modern Transformers requires solving the problems around it: position, computation, memory, and modality.
AI Fundamentals: Attention Mechanisms in Transformers (Part 2) was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.