Traditional Attention, Multi-Query Attention, and FlashAttention Overview

Issues with Traditional Attention

1. Problems with Traditional Attention

• There are constraints on context length.
• It is slow and consumes a lot of memory.

Optimization Directions for Attention

• Increase context length.
• Speed up and reduce memory usage.

Variants of Attention

• Sparse Attention: Introduces sparse biases to reduce complexity.
• Linearized Attention: Unrolls the attention matrix and kernel feature maps, achieving linear complexity through reverse order computation.
• Prototype and Memory Compression: Reduces the number of query or key-value memory pairs, shrinking the attention matrix.
• Low-order Self-Attention: Captures low-order properties of self-attention.
• Attention with Priors: Uses prior attention distributions to supplement or replace standard attention.
• Improved Multi-head Mechanism: Explores different alternatives to the multi-head mechanism.

Multi-Query Attention Overview

1. Problems with Multi-head Attention

• Training Process: Has little impact on training speed, but may result in slight model performance loss.
• Inference Process: Repeatedly loading a large KV cache incurs high memory overhead, limiting performance by memory constraints.

Introduction to Multi-Query Attention: Shares keys and values across all attention heads.

Comparison of Multi-head Attention and Multi-Query Attention

• Multi-head Attention: Each attention head has its own query, key, and value.
• Multi-query Attention: Shares keys and values across all attention heads.

Benefits of Multi-Query Attention: Reduces the size of the KV cache, decreases memory usage, and improves inference speed.

Models Using Multi-Query Attention: Representative models include PaLM, ChatGLM2, Falcon, etc.

Grouped-query Attention Overview

1. Introduction to Grouped-query Attention: Sits between multi-head and multi-query, having multiple keys and values.
2. Large Models Using Grouped-query Attention: ChatGLM2, LLaMA2-34B/70B utilize this mechanism.

FlashAttention Overview

1. Core Concept: Uses block softmax as an equivalent replacement for traditional softmax.
2. Advantages: Saves HBM, efficiently utilizes SRAM, reduces memory usage, and increases speed.
3. Representative Models: Meta’s open-source large model LLaMA and the UAE’s open-source large model Falcon use it to accelerate computation and save memory.
4. Keywords: HBM, SRAM, Block Softmax, Re-computation, Kernel Fusion.

Parallel Transformer Block Overview

Replaces serial formulas with parallel ones, achieving a 15% training speed increase for 8B parameter models with slight model performance loss, while 62B parameter models show no performance loss. Falcon and PaLM use this technology to accelerate training.