Model Architecture Design for Modern Hardware with Tri Dao

The Shift to Intelligence Per Dollar

• 💡 AI progress has been driven by scaling laws, but the focus is now shifting from simply making things work to achieving intelligence per dollar.
• 🚀 This new metric can be broken down into intelligence per flop (algorithmic efficiency) and flops per dollar (hardware efficiency).
• 🧠 Designing hardware-aware algorithms is crucial for optimizing both sides of this equation, especially by exploiting the memory hierarchy of modern hardware.

Comparing Sequence Model Architectures

• 🔄 Recurrent Neural Networks (RNNs) and State Space Models (SSMs) use a fixed-size hidden state, enabling efficient linear or constant-time inference and training.
• ⚠️ However, RNNs/SSMs struggle with information-dense modalities like language due to the strong compression required by their fixed state.
• 🧠 Transformers store the entire history (KV cache), excelling at long-range dependencies and achieving strong performance, but incur quadratic training and linear inference costs in sequence length.
• ✅ The goal is to combine the strong performance of Transformers with the efficiency and linear scaling of SSMs.

Key Ingredients for Efficient Architectures

• 📈 Large state size: Expanding the hidden state (10-100x larger than input) allows models to store significantly more information.
• 💡 Expressive state updates: Making update matrices (A, B, C) dependent on the input enables the model to selectively commit information to memory, filtering out less relevant data.
• 🛠️ Hardware-aware algorithms: This involves parallelizing computations (e.g., associative scan) and exploiting the memory hierarchy (using fast SRAM for states and HBM for inputs) and matrix multiplication units (Tensor Cores) for optimal efficiency.

Inference-First Design and Hardware Optimization

• 🚀 Inference efficiency now drives AI progress, necessitating an "inference-first" approach to model design.
• 📊 Arithmetic intensity (flops per byte loaded) is a key metric for hardware efficiency, as shown by roofline analysis.
• ✨ Techniques like Group Query Attention (GQA) and Multi-Query Attention (MQA) optimize attention for decode by increasing arithmetic intensity, leading to significantly better inference throughput.

Hybrid Models and Future Directions

• 🧩 Hybrid models, combining SSM/Mamba layers with a few attention layers (e.g., 10%), offer a **