Skip to content

Add FP16/BF16 fused_attention support with max_seqlen=512 - #175

Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged

Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom
zlsh80826:rewang/fused_attention_max_512

Conversation

@zlsh80826

@zlsh80826zlsh80826 commented Apr 25, 2023•
edited
Loading

Copy link
Copy Markdown
Collaborator

This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.

Tasks:

  • Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  • Setup the JAX unittests for the fused_attention
  • Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  • Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  • Update the documents
  • Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  • Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  • Automatically enable fused attention if available

Future works:

  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module

@zlsh80826zlsh80826 changed the title Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed the rewang/fused_attention_max_512 branch from bd1c763 to b375bd8CompareApril 27, 2023 09:07
@zlsh80826

zlsh80826 commented Apr 27, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed the rewang/fused_attention_max_512 branch from d1adeaa to 5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa reviewed May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu Outdated
Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h Outdated
@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed the rewang/fused_attention_max_512 branch from b5aec2c to df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10 reviewed May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu Outdated
@timmoon10
timmoon10 self-requested a review May 5, 2023 22:03
timmoon10
timmoon10 approved these changes May 5, 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 LGTM. The pending suggestions are mostly stylistic.

@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826

zlsh80826 commented May 8, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

ksivaman
ksivaman reviewed May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type

@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated

@zlsh80826zlsh80826May 9, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Annotation added in bd8759a

ksivaman
ksivaman reviewed May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826 requested a review from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman approved these changes May 9, 2023
View reviewed changes

@ksivamanksivaman 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.

Looks good.

cyanguwa
cyanguwa approved these changes May 9, 2023
View reviewed changes

@cyanguwacyanguwa left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Looks good to me

cyanguwa
cyanguwa approved these changes May 9, 2023
View reviewed changes

@cyanguwacyanguwa left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Looks good to me

@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826

zlsh80826 commented May 9, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@ptrendx
ptrendx merged commit 73c9f42 into NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit to nzmora-nvidia/TransformerEngine that referenced this pull request May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10 mentioned this pull request May 11, 2023
Closed
@cyanguwacyanguwa mentioned this pull request May 13, 2023
Closed
@cyanguwacyanguwa mentioned this pull request May 22, 2023
Merged
5 tasks
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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 FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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 FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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 FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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 FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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 FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@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); }
})();
})();
Add FP16/BF16 fused_attention support with max_seqlen=512 by zlsh80826 · Pull Request #175 · NVIDIA/TransformerEngine · GitHub
Skip to content
Add FP16/BF16 fused_attention support with max_seqlen=512 - #175
Merged
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
May 9, 2023
Merged
Add FP16/BF16 fused_attention support with max_seqlen=512#175
ptrendx merged 53 commits into
NVIDIA:mainfrom 
zlsh80826:rewang/fused_attention_max_512
Conversation
@zlsh80826
@zlsh80826zlsh80826
commented
Apr 25, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
This PR extends the support matrices of the following APIs: nvte_fused_attn_fwd_qkvpacked, nvte_fused_attn_bwd_qkvpacked, nvte_fused_attn_fwd_kvpacked, nvte_fused_attn_bwd_kvpacked to have the 16 bits fused attention with max sequence length = 512. Besides, this PR also adds the JAX sides custom_calls (used internally in TE) for self/cross fused attention. The TransformerLayer and MultiHeadAttention module now will automatically enable the fused attention if available.
Tasks:
  •  Rebase [WIP] Add cudnn fused multi-head attention for JAX #105 onto the main
  •  Setup the JAX unittests for the fused_attention
  •  Add fused attention 16 bits cuDNN frontend implementations under nvte_fused_attn_{fwd,bwd}_{qkv,kv}packed. (Limitation: padded shape and max sequence length = 512). cuDNN 8.9.1 is required.
  •  Merge Move dbias from fused attention bwd's input list to its output list #185 to update the latest dBias I/O
  •  Update the documents
  •  Fix a bug of generateMatrixStrides and move *_op_create to the fused_attn_utils.cu
  •  Add JAX customcalls SelfFusedAttnMax512FwdPrimitive, SelfFusedAttnMax512BwdPrimitive, CrossFusedAttnMax512FwdPrimitive, CrossFusedAttnMax512BwdPrimitive and the related custom vjp
  •  Automatically enable fused attention if available
Future works:
  • Support dropout for 16 bits fused attention
  • Add JAX DotProduct module
@zlsh80826zlsh80826
changed the title
Add fused_attention FP16/BF16 support with max_seqlen=512[WIP] Add fused_attention FP16/BF16 support with max_seqlen=512Apr 25, 2023
@zlsh80826
Add fused attention unit tests
ac39d8a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_* enums
8505428
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Use NVTE_Mask_Type and remove FMHADescriptor
f5edd9e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move common functions to utils
f2c8295
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Change namespace to fused_attn
ac07862
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_fwd_qkvpacked under the general APIs
3e33cd1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add fused_attn_max_512_bwd_qkvpacked
e4bbcd5
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Move fused_attn_max_512_bwd_qkvpacked under the general APIs
0bbaa9c
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove redundant blank line
d96c383
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Fix a potential bug for cu_seqlen converter
3fa8607
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Reformat fused_attn_max_512
b375bd8
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
bd1c763 to
b375bd8CompareApril 27, 2023 09:07
@zlsh80826
zlsh80826
commented
Apr 27, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@zlsh80826
Refine the unfused attention warning message
45735a6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_max_512
dd4877f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove the deprecated header
d91a185
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/dl/transf…
9f1c21e
…ormerengine/transformerengine into rewang/fused_attention_max_512
@zlsh80826
Fix flax import
125a1a1
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused attn
c5601b4
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attention related mask
4ff7a85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add attn_mask_type and attn_bias_type
519cc3f
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refactor jax primitive API
5dfd4fe
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
d1adeaa to
5dfd4feCompareMay 1, 2023 18:28
@zlsh80826
Merge branch 'main' into rewang/fused_attention_max_512
57cbf10
@zlsh80826
Remove dsoftmax from the customcall
44e2e6a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add None guard for bias and dropout_rng
deb4e96
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add version guard
f39e166
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add is_fused_attn_kernel_available() to correctly dispatch the attent…
8e5c92f
…ion impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
cyanguwa
cyanguwa
reviewed
May 5, 2023
View reviewed changes
Comment threaddocs/installation.rst
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_fp8.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

Comment threadtransformer_engine/common/fused_attn/utils.h
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Outdated

@zlsh80826
Refine test_fused_attn.py
638b8b9
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Unify the cuDNN compile version
857ef2e
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Add dropout to the support matrix
d717f53
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Slightly adjust the headers
df18c29
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed
the
rewang/fused_attention_max_512
branch
from
b5aec2c to
df18c29CompareMay 5, 2023 17:16
timmoon10
timmoon10
reviewed
May 5, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_max_512.cu
Outdated

@timmoon10
timmoon10
self-requested a review
May 5, 2023 22:03
timmoon10
timmoon10
approved these changes
May 5, 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 LGTM. The pending suggestions are mostly stylistic.
@zlsh80826
Typo fix: remove redundant either
b13be07
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Consolidating fused attention requirements
aadbd85
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error…
45b0f58
… line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
1f285c9
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Remove CUDNN_FRONTEND_UNUSED
74ee0e0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 8, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py

import transformer_engine_jax
from transformer_engine_jax import DType as TEDType
from transformer_engine_jax import NVTE_Bias_Type
@ksivamanksivamanMay 8, 2023

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.

Note: This file is partially annotated
@zlsh80826zlsh80826May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Annotation added in bd8759a
ksivaman
ksivaman
reviewed
May 8, 2023
View reviewed changes
Comment threadtransformer_engine/jax/cpp_extensions.py
@zlsh80826
zlsh80826
requested a review
from cyanguwaMay 8, 2023 15:57
ksivaman
ksivaman
approved these changes
May 9, 2023
View reviewed changes
@ksivamanksivaman
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.

Looks good.
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
cyanguwa
cyanguwa
approved these changes
May 9, 2023
View reviewed changes
@cyanguwacyanguwa
left a comment


Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

Looks good to me
@zlsh80826
Add more annotations to the custom calls
bd8759a
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826
commented
May 9, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@ptrendx
ptrendx
merged commit 73c9f42
into
NVIDIA:mainMay 9, 2023
nzmora-nvidia pushed a commit
to nzmora-nvidia/TransformerEngine
that referenced
this pull request
May 10, 2023
@zlsh80826@nzmora-nvidia
Add FP16/BF16 fused_attention support with max_seqlen=512 (NVIDIA#175)
8bcc11f
* Add fused attention unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_* enums
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Use NVTE_Mask_Type and remove FMHADescriptor
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move common functions to utils
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change namespace to fused_attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_fwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused_attn_max_512_bwd_qkvpacked
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Move fused_attn_max_512_bwd_qkvpacked under the general APIs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove redundant blank line
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix a potential bug for cu_seqlen converter
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Reformat fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the unfused attention warning message
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_max_512
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove the deprecated header
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix flax import
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attention related mask
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add attn_mask_type and attn_bias_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refactor jax primitive API
* Merge q_cu_seqlen and kv_cu_seqlen
* Remove is_causal_masking
* Replace seed with rng_state
* Add is_training argument
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove dsoftmax from the customcall
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add None guard for bias and dropout_rng
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add version guard
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add is_fused_attn_kernel_available() to correctly dispatch the attention impl
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the merge conflict
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Adjust the code style
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add the missing blank lines
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the order of FADescriptor members
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the readability of fused_attn_max_512.cu
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Generalize the input dimension unpacking
Signed-off-by: Reese Wang <rewang@nvidia.com>
* 16 bits fused attention requires 8.9.1
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update fused attention support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle None type when sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change to the padding ratio
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Performance optimization for non-bias cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert the cudnn-frontend PRIVATE keyword which was used for debugging
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert "Update fused attention support matrix"
This reverts commit 4effe67.
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Treat b * s as total_seqs to align ragged cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add FP16/BF16 max_seqlen <= 512 fused attention to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine test_fused_attn.py
* Replace reference code with flax.linen
* Remove unnecessary comments
* Use AttnMaskType
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Unify the cuDNN compile version
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add dropout to the support matrix
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Slightly adjust the headers
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Typo fix: remove redundant either
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Consolidating fused attention requirements
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace cudnn_frontend::throw_if with NVTE_CHECK for the better error line report
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename to fused_attn_fp16_bf16_max_seqlen_512 for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Remove CUDNN_FRONTEND_UNUSED
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add more annotations to the custom calls
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@timmoon10timmoon10
mentioned this pull request
May 11, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 13, 2023
 Closed

@cyanguwacyanguwa
mentioned this pull request
May 22, 2023
 Merged

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

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ptrendxptrendxptrendx approved these changes
@cyanguwacyanguwacyanguwa approved these changes
@ksivamanksivamanksivaman approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

5 participants

@zlsh80826@timmoon10@ptrendx@cyanguwa@ksivaman