vllm-project/vllm

[Performance]: Update Cascade Attention Heuristics for FA3

Open

#15,647 opened on Mar 27, 2025

View on GitHub
 (19 comments) (0 reactions) (1 assignee)Python (80,034 stars) (16,816 forks)batch import
feature requestgood first issuehelp wantedstaleunstale

Description

🚀 The feature, motivation and pitch

Currently, we use a heuristic (https://github.com/vllm-project/vllm/blob/4098b72210dc10761bb348b373bbd0fc9b23b0e4/vllm/v1/attention/backends/flash_attn.py#L331) to determine whether using cascade attention would improve performance. However, this heuristic was developed prior to the FA3 integration and is therefore optimized only for FA2. The calculation of SM occupancy it uses is no longer accurate for FA3 and needs updating.

Specifically,

  1. We need to split the case for FA2 and FA3, since FA2 is still used for certain GPUs.
  2. Afaik, FA3 uses different heuristics for GQA packing and split kv than FA2. The heuristics in use_cascade should reflect this difference (although it doesn't need to be super accurate).
  3. It'd be nice if we can cite specific lines of code in FA3 (and FA2) deciding the tile sizes, scheduling, etc., so that we can easily verify and track.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Contributor guide

Tech stack
python
Domain
performance
Issue type
performance
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
half day
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
blocked
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
Familiarity with Flash Attention (FA2/FA3)Understanding of GPU occupancy heuristicsKnowledge of vllm codebase structure
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
25
Research direction
Examine the current heuristic in vllm/v1/attention/backends/flash attn.py at line 331. Split the logic for FA2 and FA3, referencing FA3's internal heuristics for tile sizes and scheduling (e.g., in Flash Attention repository or vllm's FA3 integration). Update the function `use cascade` accordingly, and add comments citing specific lines of FA3 code that determine tile sizes and scheduling decisions.