Exact Causal Attention with 10% Fewer Operations

2025-10-05

AI Summary

This paper addresses the high computational cost of triangular matrix multiplication (e.g., Mask(QKᵀ)) in causal attention. We propose Fast Causal Attention (FCA), the first approach to reduce floating-point operations by ~10% *without sacrificing numerical precision*, achieved by discovering novel algebraic identities via machine learning–driven combinatorial search and restructuring the causal attention computation graph. Based on these identities, FCA designs a GPU-optimized, production-ready kernel compatible with both PyTorch and Triton, significantly accelerating forward and backward passes. Experiments demonstrate that FCA outperforms PyTorch’s native implementation and Triton-compiled kernels across mainstream models, yielding substantial training latency reduction. Our core contributions are: (i) the first numerically exact algebraic optimization framework for causal attention; and (ii) a deployable, hardware-efficient kernel implementation that bridges algorithmic innovation with practical system performance.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

We present Fast Causal Attention (FCA), an algorithm that computes exact Causal Attention using 10% fewer operations. FCA accelerates a special class of matrix multiplications where either one operand or the output matrix is upper- or lower-triangular. This includes all operations in forward and backward pass of Causal Attention, such as masked product $mathrm{Mask}(QK^{T})$. For these matrix multiplications on GPU, FCA reaches noticeable accelerations over the default PyTorch implementations and Triton compiled kernels. FCA is built upon algebraic identities discovered via machine learning and combinatorial search.

Problem

• Reduces computational operations in exact causal attention
• Accelerates triangular matrix multiplications for causal attention
• Improves GPU performance over PyTorch and Triton implementations

Innovation

• Fast Causal Attention reduces operations by 10%
• Accelerates triangular matrix multiplications on GPU
• Uses machine learning discovered algebraic identities