Skip to content

[PyTorch] Support for cuDNN-backed flex attention #2984

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vcherepanov-nv merged 30 commits into from
Jun 4, 2026
Merged

[PyTorch] Support for cuDNN-backed flex attention #2984

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
11c3ed2
Add cuDNN score_mod attention path
vcherepanov-nv May 8, 2026
eb35191
Avoid BHSD copies in score_mod attention
vcherepanov-nv May 11, 2026
57ce106
Test relative position score_mod attention
vcherepanov-nv May 11, 2026
e6ba0ea
Test softcap score_mod attention
vcherepanov-nv May 11, 2026
dcb6b49
Run score_mod graphs on current CUDA stream
vcherepanov-nv May 11, 2026
fefcbe7
Add PyTorch score_mod execution plan cache
vcherepanov-nv May 12, 2026
ac4c60d
Fix score_mod cache edge cases
vcherepanov-nv May 12, 2026
6446825
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 13, 2026
58a5fb5
Fix score_mod callback graph cache keys
vcherepanov-nv May 15, 2026
c00a0b7
Address score_mod review feedback
vcherepanov-nv May 15, 2026
a8ed67e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 15, 2026
e2a69e1
Fix score_mod lambda cache keys
vcherepanov-nv May 15, 2026
96f8ab2
Address flex attention review feedback
vcherepanov-nv May 19, 2026
e11cc23
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 19, 2026
f0f4f7b
Address flex attention backend review feedback
vcherepanov-nv May 21, 2026
9420479
Validate score_mod bprop tensor inputs
vcherepanov-nv May 21, 2026
db4e7c6
Add flex-attn tests to QA scripts
vcherepanov-nv May 22, 2026
6ca49cf
Add lint directive for score_mod tensor type checks
vcherepanov-nv May 22, 2026
a93df53
Alias dataclass field import in attention utils
vcherepanov-nv May 22, 2026
45ae2db
Skip softcap flex attention tests before sm90
vcherepanov-nv May 23, 2026
06605aa
Address flex attention review feedback
vcherepanov-nv Jun 2, 2026
9cbab29
Address attention backend review nits
vcherepanov-nv Jun 2, 2026
297e66e
Remove duplicate flex attention asserts
vcherepanov-nv Jun 2, 2026
7f514e1
Clarify score mod tensor keys
vcherepanov-nv Jun 2, 2026
fb51502
Normalize Flex Attention naming
vcherepanov-nv Jun 2, 2026
b82830e
Simplify score mod backward graph cache lookup
vcherepanov-nv Jun 2, 2026
4deeb2a
Return only cuDNN graph from helper
vcherepanov-nv Jun 2, 2026
e9fcf7c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 2, 2026
2e1874e
Refer to Flex Attention in error messages
vcherepanov-nv Jun 3, 2026
984cf47
Address Flex Attention review comments
vcherepanov-nv Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions qa/L0_pytorch_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_e
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py"
NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py"
Expand Down
Loading
Loading