K.Boopathi
BlogNotesSecond BrainProjectsHire MeYour reading file0 pts

Attention And Types

Oct 03, 20253 min read

  • **Multi-Head Attention (MHA)
  • Flash Attention
  • Group Query Attention
  • Mixture of Experts
  • Multi-Query Attention (MQA)
  • Grouped-Query Attention (GQA)
  1. GPT-4 and Similar Models: Primarily use optimized multi-head attention with techniques like Flash Attention for efficiency
  2. Llama and Mistral Family: Implement Grouped-Query Attention for improved inference efficiency while maintaining model quality
  3. Longformer and Long-Context Models: Employ sliding window attention combined with global attention for handling extended sequences
  4. Mixture of Experts (MoE) Models: Often combine GQA with sparse architectures for both attention and feed-forward efficiency

Ring Attention

For distributing the model to multiple GPU https://coconut-mode.com/posts/ring-attention/

Flash Attention

It is fundamentally an IO-aware algorithm that optimizes how data moves within a GPU to prevent memory transfer from becoming a bottleneck.

The Core Problem: The Memory Wall Standard attention requires massive amounts of read/write operations between the GPU’s two memory types:

  • HBM (High Bandwidth Memory): Large capacity but relatively slow (tens of gigabytes) .
  • SRAM (Static Random Access Memory): Very fast but tiny capacity, located directly on-chip next to the compute units (tens of megabytes) .

In standard attention, the softmax operation forces the GPU to write large intermediate matrices back to HBM before the next step can begin, creating a massive data transfer bottleneck that slows down the entire training process .

The Flash Attention Solution

Flash Attention introduces two main strategies to solve this:

  • Tiling : Instead of processing the entire attention matrix at once, the algorithm breaks the input matrices into small blocks (tiles). These blocks are loaded into the SRAM, where the entire attention computation (including the softmax) is completed locally. The results are only written back to HBM once the block is fully processed. This drastically reduces the number of read/write cycles to the slow HBM.

  • Recomputation : During the backwards pass of training, instead of storing all intermediate activations in HBM (which consumes huge amounts of memory), Flash Attention recomputes the activations on the fly during the backward pass using the fast SRAM. Because the attention operation is now so much faster, it is actually more efficient to recompute these values than to store them in memory and reload them later.

Impact

  • Memory Efficiency: It significantly lowers the memory footprint by avoiding the storage of large intermediate matrices .

  • Speed: By eliminating the memory bottleneck, it achieves substantial speedups compared to standard attention, allowing for longer context windows and faster training .

  • Exactness: Crucially, this is an exact computation; it does not approximate the attention scores, meaning it provides the same mathematical result as standard attention while being significantly faster .

Graph View