Skip to content

[PyTorch] Add sink attention support from cuDNN #2148

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
merged 60 commits into from
Sep 22, 2025
Merged

[PyTorch] Add sink attention support from cuDNN #2148

Show file tree
Hide file tree
Changes from 58 commits
Commits
Show all changes
60 commits
Select commit Hold shift + click to select a range
26e1deb
first draft; debug plan failure
cyanguwa Jul 23, 2025
73f2ad3
debug uid error
cyanguwa Jul 23, 2025
c3c1843
tweak params
cyanguwa Jul 23, 2025
6e59c49
add grad in output
cyanguwa Jul 23, 2025
a2242e8
clean up prints
cyanguwa Jul 23, 2025
854cf1f
fix prints in test
cyanguwa Jul 23, 2025
95f44fc
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Jul 24, 2025
308c3aa
address review comments
cyanguwa Jul 25, 2025
4050332
fix unfused grad; add softmax_type; add sink to bwd
cyanguwa Jul 29, 2025
fdbdabc
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
cfab71c
fix padding mask; add swa tests; remove requires_grad for off-by-one
cyanguwa Aug 1, 2025
b8ff061
update FE
cyanguwa Aug 1, 2025
9aa99c1
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
cde079e
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
be47d64
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
ed0b389
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
802b552
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
d879a77
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
f48b5fc
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
557f982
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
38c4cb6
Apply 1 suggestion(s) to 1 file(s)
cyanguwa Aug 1, 2025
094f9f8
fix indent
cyanguwa Aug 4, 2025
2e8acc5
fix non-determinism and shapes
cyanguwa Aug 4, 2025
13b8d99
clean up prints
cyanguwa Aug 4, 2025
a44fe19
add GQA
cyanguwa Aug 4, 2025
14122ff
add CP A2A; dq/dk mismatches
cyanguwa Aug 6, 2025
0341c49
fix CP A2A; need cleaner solution
cyanguwa Aug 6, 2025
d7053fe
fix CP A2A; pending cudnn kernel change
cyanguwa Aug 6, 2025
c259b3c
minor fixes
cyanguwa Aug 7, 2025
31cdaf9
fix world size in unit test; avoid thd format
cyanguwa Aug 7, 2025
b934d1c
fix kernel_backend, dtype in unit test; fix head_dim for FP8 Hopper
cyanguwa Aug 7, 2025
24b7288
fix thd logic
cyanguwa Aug 7, 2025
ffdb634
fix fp8 context
cyanguwa Aug 7, 2025
08bc106
tweak CP logging
cyanguwa Aug 7, 2025
7ef12bf
allow no_mask/padding for SWA(left,0)
cyanguwa Aug 7, 2025
6736fa6
Revert "allow no_mask/padding for SWA(left,0)"
cyanguwa Aug 7, 2025
be8e3a9
add softmax_type to Jax
cyanguwa Aug 7, 2025
f2d57b6
add cuDNN version control
cyanguwa Aug 7, 2025
ff450fa
prettify tests
cyanguwa Aug 8, 2025
17db40b
skip 9.13 for MLA, non 192/128
cyanguwa Aug 8, 2025
a5cf6e8
rename compare_with_error
cyanguwa Aug 8, 2025
edfd8de
small cleanups and improvements
cyanguwa Aug 8, 2025
eeb7a53
fix minor CI failures
cyanguwa Aug 9, 2025
1121461
force sink/dsink to be float32
cyanguwa Aug 28, 2025
65d6138
switch FE to GH FE
cyanguwa Sep 2, 2025
9828253
return to GH TE main FE commit
cyanguwa Sep 2, 2025
7ad9aba
Merge branch 'main' into sink_attn
cyanguwa Sep 2, 2025
f58e707
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 2, 2025
a15f079
update FE to 1.14.1
cyanguwa Sep 5, 2025
f3c17eb
Merge branch 'main' into sink_attn
cyanguwa Sep 5, 2025
11e4329
Merge branch 'main' into sink_attn
cyanguwa Sep 11, 2025
f1c1688
clean up before CI
cyanguwa Sep 11, 2025
a701b17
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 11, 2025
516a423
fix lint
cyanguwa Sep 11, 2025
435c338
bump up cudnn version
cyanguwa Sep 11, 2025
0822c81
Merge branch 'main' into sink_attn
cyanguwa Sep 17, 2025
e97327d
Merge branch 'main' into sink_attn
cyanguwa Sep 18, 2025
26adddb
Merge branch 'main' into sink_attn
cyanguwa Sep 20, 2025
d8a7d6f
add backend selection guard for unit tests
cyanguwa Sep 21, 2025
759b5ab
add docstring for softmax type enums in C
cyanguwa Sep 22, 2025
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
2 changes: 1 addition & 1 deletion 3rdparty/cudnn-frontend
Open in desktop
Submodule cudnn-frontend updated 18 files
+1 −1 CMakeLists.txt
+2 −1 README.md
+1 −0 dlpack_version.txt
+54 −3 include/cudnn_frontend/graph_properties.h
+12 −0 include/cudnn_frontend/node/pointwise.h
+43 −6 include/cudnn_frontend/node/scaled_dot_product_flash_attention.h
+0 −560 include/cudnn_frontend/node/sdpa_fp8.h
+32 −4 include/cudnn_frontend/node/softmax.h
+1 −1 include/cudnn_frontend_version.h
+5 −1 python/CMakeLists.txt
+1 −1 python/cudnn/__init__.py
+41 −2 python/pygraph/pointwise.cpp
+13 −0 python/pygraph/pygraph.h
+2 −0 samples/cpp/CMakeLists.txt
+332 −0 samples/cpp/sdpa/fp16_bwd_with_sink_token.cpp
+254 −0 samples/cpp/sdpa/fp16_fwd_with_sink_token.cpp
+5 −0 setup.py
+1 −1 test/python/test_matmul_bias_relu.py
Loading