Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sawtooth Wavefront Reordering: Enhanced CuTile FlashAttention on NVIDIA GB10

Yifan Zhu, Yekai Pan, Chen Ding

TL;DR

The paper tackles the cache bottlenecks of high-performance CuTile Flash Attention kernels on NVIDIA GB10, identifying L2 non-compulsory misses as the main limiter for large KV workloads. It combines hardware-counter–driven analysis with a deterministic CTAs scheduling study to reveal that L1 caching provides little benefit for streaming attention, and that L2 reuse scales with the number of active SMs, suggesting wavefront reuse as a key factor. The authors introduce Sawtooth Wavefront Reordering, which alternates KV traversal to reduce reuse distance and thereby dramatically cut L2 misses, validated in both CUDA and CuTile environments with substantial throughput gains (roughly 50%–60% depending on variant). A lightweight L2 sector-access model is developed to predict memory traffic and to explain observed improvements, and the approach generalizes to tile-based programming models. These results advance cache-aware design for large-sequence attention and offer practical guidance for achieving higher throughput on modern GPUs with complex memory hierarchies.

Abstract

High-performance attention kernels are essential for Large Language Models. This paper presents analysis of CuTile-based Flash Attention memory behavior and a technique to improve its cache performance. In particular, our analysis on the NVIDIA GB10 (Grace Blackwell) identifies the main cause of L2 cache miss. Leveraging this insight, we introduce a new programming technique called Sawtooth Wavefront Reordering that reduces L2 misses. We validate it in both CUDA and CuTile, observing 50\% or greater reduction in L2 misses and up to 60\% increase in throughput on GB10.

Sawtooth Wavefront Reordering: Enhanced CuTile FlashAttention on NVIDIA GB10

TL;DR

The paper tackles the cache bottlenecks of high-performance CuTile Flash Attention kernels on NVIDIA GB10, identifying L2 non-compulsory misses as the main limiter for large KV workloads. It combines hardware-counter–driven analysis with a deterministic CTAs scheduling study to reveal that L1 caching provides little benefit for streaming attention, and that L2 reuse scales with the number of active SMs, suggesting wavefront reuse as a key factor. The authors introduce Sawtooth Wavefront Reordering, which alternates KV traversal to reduce reuse distance and thereby dramatically cut L2 misses, validated in both CUDA and CuTile environments with substantial throughput gains (roughly 50%–60% depending on variant). A lightweight L2 sector-access model is developed to predict memory traffic and to explain observed improvements, and the approach generalizes to tile-based programming models. These results advance cache-aware design for large-sequence attention and offer practical guidance for achieving higher throughput on modern GPUs with complex memory hierarchies.

Abstract

High-performance attention kernels are essential for Large Language Models. This paper presents analysis of CuTile-based Flash Attention memory behavior and a technique to improve its cache performance. In particular, our analysis on the NVIDIA GB10 (Grace Blackwell) identifies the main cause of L2 cache miss. Leveraging this insight, we introduce a new programming technique called Sawtooth Wavefront Reordering that reduces L2 misses. We validate it in both CUDA and CuTile, observing 50\% or greater reduction in L2 misses and up to 60\% increase in throughput on GB10.
Paper Structure (21 sections, 3 equations, 12 figures, 3 tables, 4 algorithms)

This paper contains 21 sections, 3 equations, 12 figures, 3 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. GPU Memory Hierarchy
  4. NVIDIA GB10
  5. Nsight Compute CLI
  6. Flash Attention and Implmementation
  7. Tiling Implementation
  8. CTA Scheduling and Work Distribution
  9. Cache Performance Analysis
  10. Effect of L1 Caching
  11. Modeling L2 Sector Access
  12. L2 Non-Compulsory Miss Threshold
  13. The Cause of L2 Non-Compulsory Miss
  14. Sawtooth Wavefront Reordering
  15. Implementation
  16. ...and 6 more sections

Figures (12)

  • Figure 1: L1/L2 Metrics for Sequence Length 32K. Parameters: Batch=1, Heads=1, Head Dim=64, Tile Size=80x80.
  • Figure 2: L1/L2 Metrics for Sequence Length 128K. Parameters: Batch=1, Heads=1, Head Dim=64, Tile Size=80x80.
  • Figure 3: L2 Sector Access vs Sequence Length (Non-Causal Masking, $T=80$).
  • Figure 4: L2 Sector Access vs Sequence Length (Causal Masking, $T=80$).
  • Figure 5: L2 Miss Count vs Sequence Length (SM=48). The dashed line represents the theoretical cold miss count ($16S$).
  • ...and 7 more figures