Skip to content

Fix Flash and Lean attention split heuristics - #29554

Merged
Tianlei Wu (tianleiwu) merged 2 commits into
microsoft:mainfrom
tianleiwu:tlwu/fix_flash_lean_num_splits_29550
Jul 6, 2026
Merged

Tianlei Wu (tianleiwu) merged 2 commits into
microsoft:mainfrom
tianleiwu:tlwu/fix_flash_lean_num_splits_29550

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

Description

Fixes split heuristic edge cases in Flash Attention and Lean Attention when the SM count is reported as zero or when there are no key tiles. The change clamps the SM count used by the heuristics, prevents divide-by-zero paths, and returns stable buffer sizing for empty tile cases.

Changes

  • Guard Flash Attention split heuristic against zero SM count and degenerate split counts.
  • Guard Lean Attention split/buffer sizing against zero SM count, zero key tiles, and single-tile division cases.
  • Add CUDA provider regression tests for zero SM count and zero key-tile inputs.

Testing

  • Added attention_split_heuristic_test.cc coverage for Flash Attention and Lean Attention split heuristics.

Fixes #29550

@tianleiwu
Tianlei Wu (tianleiwu) merged commit 4942b12 into microsoft:main Jul 6, 2026
87 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[CUDA] Attention split heuristics unsafe at small SM counts (MPS partitions): reachable div-by-zero SIGFPE in lean_attention, degenerate flash splits

2 participants