Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The Anatomy of a Triton Attention Kernel

Burkhard Ringlein, Jan van Lunteren, Radu Stoica, Thomas Parnell

TL;DR

This paper tackles the challenge of portable, high-performance LLM inference across GPU vendors by developing a Triton-based paged attention kernel. It demonstrates state-of-the-art performance on both NVIDIA and AMD GPUs and integrates the solution into the open-source vLLM inference server, making it broadly accessible. The authors detail kernel designs (including Q-Block, tiled softmax, and static launch grids), autotuning heuristics, and system integration, showing end-to-end latency improvements up to $5.9\times$ and near parity with Flash Attention on major hardware. The work highlights how open-source DSLs like Triton can unlock cross-vendor portability and practical deployment of cutting-edge attention kernels, with open-source releases and clear directions for future GPU-specific tuning and graph usage.

Abstract

A long-standing goal in both industry and academia is to develop an LLM inference platform that is portable across hardware architectures, eliminates the need for low-level hand-tuning, and still delivers best-in-class efficiency. In this work, we demonstrate that portable, efficient cross-platform LLM inference is indeed possible and share our experience. We develop a state-of-the-art paged attention kernel, the core performance-critical component of many LLM deployments, that builds exclusively on the domain-specific just-in-time compiled language Triton to achieve state-of-the-art performance on both NVIDIA and AMD GPUs. We describe our high-level approach, the key algorithmic and system-level improvements, the parameter auto-tuning required to unlock efficiency, and the integrations into a popular inference server that are necessary to bring the performance of a generic Triton attention kernel from 19.7% of the state-of-the-art to 105.9%. Our results highlight how open-source domain-specific languages can be leveraged to unlock model portability across different GPU vendors.

The Anatomy of a Triton Attention Kernel

TL;DR

This paper tackles the challenge of portable, high-performance LLM inference across GPU vendors by developing a Triton-based paged attention kernel. It demonstrates state-of-the-art performance on both NVIDIA and AMD GPUs and integrates the solution into the open-source vLLM inference server, making it broadly accessible. The authors detail kernel designs (including Q-Block, tiled softmax, and static launch grids), autotuning heuristics, and system integration, showing end-to-end latency improvements up to and near parity with Flash Attention on major hardware. The work highlights how open-source DSLs like Triton can unlock cross-vendor portability and practical deployment of cutting-edge attention kernels, with open-source releases and clear directions for future GPU-specific tuning and graph usage.

Abstract

A long-standing goal in both industry and academia is to develop an LLM inference platform that is portable across hardware architectures, eliminates the need for low-level hand-tuning, and still delivers best-in-class efficiency. In this work, we demonstrate that portable, efficient cross-platform LLM inference is indeed possible and share our experience. We develop a state-of-the-art paged attention kernel, the core performance-critical component of many LLM deployments, that builds exclusively on the domain-specific just-in-time compiled language Triton to achieve state-of-the-art performance on both NVIDIA and AMD GPUs. We describe our high-level approach, the key algorithmic and system-level improvements, the parameter auto-tuning required to unlock efficiency, and the integrations into a popular inference server that are necessary to bring the performance of a generic Triton attention kernel from 19.7% of the state-of-the-art to 105.9%. Our results highlight how open-source domain-specific languages can be leveraged to unlock model portability across different GPU vendors.

Paper Structure

This paper contains 36 sections, 2 equations, 9 figures.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. vLLM: An Inference Server
  4. Triton: A tiling DSL
  5. Portability and Kernel Parameter Autotuning
  6. Attention Kernels and Libraries
  7. Solution Overview
  8. Design and Implementation of Attention
  9. Core Concepts
  10. Terminology
  11. Baseline Triton Attention Kernel
  12. Prefill and GQA Optimization
  13. Parallel Tiled Softmax
  14. Adjustable Tile Sizes
  15. Static Launch Grid
  16. ...and 21 more sections

Figures (9)

  • Figure 1: Architecture of vLLM v1 and the question of cross-vendor performance portability.
  • Figure 2: Overview of the triton_attn backend in vLLM.
  • Figure 3: Q Block and KV Tiles.
  • Figure 4: Parallel Tiled Softmax.
  • Figure 5: Workflow of tuning in vllm-triton-backend.
  • ...and 4 more figures