Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sequence-Aware Split Heuristic to Mitigate SM Underutilization in FlashAttention-3 Low-Head-Count Decoding

Martí Llopart Font, Javier Hernando, Cristina España-Bonet

Abstract

The standard FlashAttention-3 heuristic exhibits a GPU occupancy bottleneck in low-head-count decoding configurations because it disables sequence splitting based on sequence length alone, underutilizing the Streaming Multiprocessors of Hopper GPUs. Our proposed sequence-aware split policy mitigates this by allowing sequence-level parallelism in low-head-count regimes, improving hardware utilization to deliver roughly a 21 to 24% improvement in decoder kernel efficiency on metadata-enabled inference paths, with no observed regressions.

Sequence-Aware Split Heuristic to Mitigate SM Underutilization in FlashAttention-3 Low-Head-Count Decoding

Abstract

The standard FlashAttention-3 heuristic exhibits a GPU occupancy bottleneck in low-head-count decoding configurations because it disables sequence splitting based on sequence length alone, underutilizing the Streaming Multiprocessors of Hopper GPUs. Our proposed sequence-aware split policy mitigates this by allowing sequence-level parallelism in low-head-count regimes, improving hardware utilization to deliver roughly a 21 to 24% improvement in decoder kernel efficiency on metadata-enabled inference paths, with no observed regressions.

Paper Structure

This paper contains 17 sections, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. Occupancy Collapse in Low-Head-Count Decoding
  4. The Premature Guard Flaw
  5. Automated Discovery via Evolutionary Search
  6. Experimental Design and Search Space
  7. Evolutionary Process
  8. Empirical Dissection and Motivation
  9. Sequence-Aware Split Heuristic
  10. Scope: Why L_K=512 and Not Shorter?
  11. The C++ Heuristic Patch
  12. Empirical Evaluation
  13. Kernel-Level Speedups (A/B Testing)
  14. Extended Split Sweep and the Chosen Split Count
  15. Safety and Regression Profiling
  16. ...and 2 more sections

Figures (3)

  • Figure 1: A fragment of the high-performing evolved Python heuristic.
  • Figure 2: Conservative C++ policy used in the paper: keep shorter and saturated cases unchanged, and override the low-tile $nblk=4$ boundary bucket with $s=3$.
  • Figure 3: Extended kernel-level split sweep for $Batch=1$, $L_K=512$, $H_{KV}=1$, and $D=128$ using precomputed scheduler metadata. Latency drops sharply once sequence splitting is enabled and then flattens into a broad low-latency plateau with shallow local minima at larger split counts. The best tested value in this sweep is $s=64$ ($\sim$11.14 ); the policy uses $s=3$ as a safeguard (smallest split entering this regime), with gain under $\sim$2%.