Skip to content

Add FP8 fused attention - #155

Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged

Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom
cyanguwa:fp8_flash_attn_new_api

Conversation

@cyanguwa

@cyanguwacyanguwa commented Apr 17, 2023•
edited
Loading

Copy link
Copy Markdown
Collaborator

This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.

  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked

It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.

timmoon10
timmoon10 reviewed Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt Outdated
Comment threadtransformer_engine/common/CMakeLists.txt Outdated
Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu Outdated
@cyanguwacyanguwa mentioned this pull request Apr 17, 2023
Closed
timmoon10
timmoon10 reviewed Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h Outdated
Comment threadtransformer_engine/common/fused_attn/utils.h Outdated
@cyanguwa
cyanguwa marked this pull request as ready for review April 17, 2023 22:27
@cyanguwa

cyanguwa commented Apr 17, 2023•
edited
Loading

Copy link
Copy Markdown
CollaboratorAuthor

Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.

timmoon10
timmoon10 approved these changes Apr 17, 2023
View reviewed changes

@timmoon10timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.

Comment threaddocs/installation.rst Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10

timmoon10 commented Apr 18, 2023

Copy link
Copy Markdown
Member

/te-ci

timmoon10
timmoon10 reviewed Apr 18, 2023
View reviewed changes
Comment threadsetup.py Outdated
@timmoon10

timmoon10 commented Apr 18, 2023

Copy link
Copy Markdown
Member

/te-ci

1 similar comment
@timmoon10

timmoon10 commented Apr 18, 2023

Copy link
Copy Markdown
Member

/te-ci

@cyanguwa

cyanguwa commented Apr 18, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te_ci

1 similar comment
@cyanguwa

cyanguwa commented Apr 18, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te_ci

@timmoon10

timmoon10 commented Apr 18, 2023

Copy link
Copy Markdown
Member

/te-ci

1 similar comment
@timmoon10

timmoon10 commented Apr 18, 2023

Copy link
Copy Markdown
Member

/te-ci

@ptrendx

ptrendx commented Apr 19, 2023

Copy link
Copy Markdown
Member

@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.

@cyanguwa

cyanguwa commented Apr 19, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed the fp8_flash_attn_new_api branch from 8384086 to c3795e0CompareApril 19, 2023 17:34
@cyanguwa

cyanguwa commented Apr 19, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10 reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h Outdated
Comment threaddocs/installation.rst Outdated
@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed the fp8_flash_attn_new_api branch from eb2e5e8 to 4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa reopened this Apr 20, 2023
@cyanguwa

cyanguwa commented Apr 20, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx reviewed Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst Outdated
ptrendx
ptrendx reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt Outdated
@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
ptrendx
ptrendx reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa

cyanguwa commented Apr 20, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Apr 21, 2023

Copy link
Copy Markdown
Member

/te-ci

ptrendx
ptrendx reviewed Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h Outdated
@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa

cyanguwa commented Apr 21, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa

cyanguwa commented Apr 21, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa

cyanguwa commented Apr 21, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@ksivaman

ksivaman commented Apr 21, 2023

Copy link
Copy Markdown
Member

Pipeline 8043078

@ksivaman
ksivaman merged commit 989a53a into NVIDIA:mainApr 21, 2023
@timmoon10timmoon10 mentioned this pull request Apr 24, 2023
Closed
@ksivamanksivaman mentioned this pull request Apr 25, 2023
Merged
ptrendx pushed a commit that referenced this pull request Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa mentioned this pull request May 13, 2023
Closed
@cyanguwacyanguwa mentioned this pull request May 22, 2023
Merged
5 tasks
@hungryGeek16hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees

No one assigned

Labels

0.8.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
Add FP8 fused attention by cyanguwa · Pull Request #155 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP8 fused attention - #155
Merged
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Apr 21, 2023
Merged
Add FP8 fused attention#155
ksivaman merged 17 commits into
NVIDIA:mainfrom 
cyanguwa:fp8_flash_attn_new_api
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR adds the cuDNN-based fused attention to Transformer Engine. It provides two low-level PyTorch and C APIs for the forward and backward passes of a dot product attention calculation, i.e. BMM1 -> Softmax -> Dropout -> BMM2.
  • pytorch.cpp_extensions: fused_attn_fwd_qkvpacked, fused_attn_bwd_qkvpacked, fused_attn_fwd_kvpacked, fused_attn_bwd_kvpacked
  • C API: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked
It focuses on FP8 and PyTorch but it provides a template for adding other future fused attention implementations to Transformer Engine.
timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/CMakeLists.txt
Outdated

Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

Comment threadtransformer_engine/pytorch/constants.py
Comment threadtransformer_engine/pytorch/csrc/extensions.cu
Outdated

@cyanguwacyanguwa
mentioned this pull request
Apr 17, 2023
 Closed

timmoon10
timmoon10
reviewed
Apr 17, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Outdated

@cyanguwa
cyanguwa
marked this pull request as ready for review
April 17, 2023 22:27
@cyanguwa
cyanguwa
commented
Apr 17, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
Currently, the cudnn-frontend submodule points to v0.9-rc of NVIDIA/cudnn-frontend. Will point it to main when PR 44 is merged.
timmoon10
timmoon10
approved these changes
Apr 17, 2023
View reviewed changes
@timmoon10timmoon10
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. I'm happy once the tests are green.
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
timmoon10
timmoon10
reviewed
Apr 18, 2023
View reviewed changes
Comment threadsetup.py
Outdated

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci

1 similar comment

@cyanguwa
cyanguwa
commented
Apr 18, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te_ci
@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci

1 similar comment

@timmoon10
timmoon10
commented
Apr 18, 2023

Copy link


Copy Markdown

Member
/te-ci
@ptrendx
ptrendx
commented
Apr 19, 2023

Copy link


Copy Markdown

Member
@cyanguwa please clean the commit history in this PR - right now it includes a whole bunch of unrelated commits.
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
Add FP8 fused attention to TE for PyTorch
c3795e0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
8384086 to
c3795e0CompareApril 19, 2023 17:34
@cyanguwa
cyanguwa
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add license for cudnn-frontend, modify installation requirements, and…
db90013
… refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
timmoon10
timmoon10
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/transformer_engine.h
Outdated

Comment threaddocs/installation.rst
Outdated

@cyanguwa
add c api docs for fused attention
dbaf22d
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
add exception for unsupported precision/sequence length combinations
eb2e5e8
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwaforce-pushed
the
fp8_flash_attn_new_api
branch
from
eb2e5e8 to
4a1efe8CompareApril 20, 2023 20:16
@cyanguwa
merge other changes on NVDIA/TE/main to fp8_flash_attn_new_api
63cdc82
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwacyanguwa
reopened this
Apr 20, 2023
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix installation requirement for non fused attn use cases
f2188f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
98f9da7
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/CMakeLists.txt
Outdated

@cyanguwa
Merge branch 'fp8_flash_attn_new_api' of github.com:cyanguwa/Transfor…
000d345
…merEngine into fp8_flash_attn_new_api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

@cyanguwa
minor fixes based on PR comments
3178c03
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix description for kvpacked fwd
2243003
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

ptrendx
ptrendx
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
@cyanguwa
fix description of Bias in C api
ec53e9a
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
minor fixes for cudnn requirement and description for QKV tensors
7bb62a5
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 20, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman@cyanguwa
Merge remote-tracking branch 'main' into fp8_flash_attn_new_api
ad80f4b
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
/te-ci
ptrendx
ptrendx
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@cyanguwa
fix QKV layout description and support matrix for C api
06cbca0
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
add asserts to cpp_extensions for qkv layout/bias type/attn mask type
40502ef
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@cyanguwa
fix typo precision
bb415f9
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
cyanguwa
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ksivaman
ksivaman
commented
Apr 21, 2023

Copy link


Copy Markdown

Member
Pipeline 8043078
@ksivaman
ksivaman
merged commit 989a53a
into
NVIDIA:mainApr 21, 2023
@timmoon10timmoon10
mentioned this pull request
Apr 24, 2023
 Closed

@ksivamanksivaman
mentioned this pull request
Apr 25, 2023
 Merged

ptrendx pushed a commit
that referenced
this pull request
Apr 25, 2023
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

5 tasks
@hungryGeek16hungryGeek16
mentioned this pull request
May 31, 2026
 Open

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ptrendxptrendxptrendx approved these changes
@timmoon10timmoon10timmoon10 approved these changes

Assignees


No one assigned


Labels

0.8.0

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@timmoon10@ptrendx@ksivaman