') + ')', '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('^' + ".*" + ', '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" + ', '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('^' + ".*" + ', '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 long sequence support for fused attention by cyanguwa · Pull Request #237 · NVIDIA/TransformerEngine · GitHub
Skip to content

Add long sequence support for fused attention - #237

Merged
ksivaman merged 27 commits into
NVIDIA:mainfrom
cyanguwa:fused_attn/fp8_f16_c_pytorch
Jun 22, 2023
Merged

Add long sequence support for fused attention#237
ksivaman merged 27 commits into
NVIDIA:mainfrom
cyanguwa:fused_attn/fp8_f16_c_pytorch

Conversation

@cyanguwa

@cyanguwacyanguwa commented May 22, 2023•
edited
Loading

Copy link
Copy Markdown
Collaborator

This PR provides further support for fused attention.

  • Add PyTorch support for BF16/FP16 fused attention with <=512 sequence length. This is based the C functionality already added by Add FP16/BF16 fused_attention support with max_seqlen=512 #175.
  • Add C and PyTorch support for BF16/FP16 fused attention with any arbitrary sequence length.
  • Add PyTorch unit tests for FP8 fused attention (Add FP8 fused attention #155) with <=512 sequence length.
  • Add selection mechnism for the three backends above, on both C and PyTorch levels.
  • Clean up naming, coding style, etc to make three backends look more uniform.

@cyanguwa
cyanguwa requested review from ptrendx and zlsh80826 and removed request for ptrendx and zlsh80826May 22, 2023 15:33
@cyanguwa
cyanguwaforce-pushed the fused_attn/fp8_f16_c_pytorch branch from 2f2e2e3 to 6900396CompareMay 23, 2023 06:38
@cyanguwacyanguwa reopened this May 23, 2023
@cyanguwacyanguwa closed this Jun 5, 2023
@cyanguwa
cyanguwaforce-pushed the fused_attn/fp8_f16_c_pytorch branch 2 times, most recently from adfc650 to 144e488CompareJune 5, 2023 23:05
@cyanguwacyanguwa reopened this Jun 5, 2023
@cyanguwacyanguwa closed this Jun 6, 2023
@cyanguwa
cyanguwaforce-pushed the fused_attn/fp8_f16_c_pytorch branch 2 times, most recently from 456b4ce to 207b231CompareJune 6, 2023 21:09
@cyanguwacyanguwa reopened this Jun 6, 2023
@cyanguwacyanguwa changed the title [WIP] Add C and PyTorch support for fused attentionAdd C and PyTorch support for fused attentionJun 12, 2023
@cyanguwacyanguwa changed the title Add C and PyTorch support for fused attentionAdd long sequence support for fused attentionJun 12, 2023
@cyanguwa

cyanguwa commented Jun 12, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
add long sequence support and unify three backends for fused attention
a78ed36
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
cyanguwaforce-pushed the fused_attn/fp8_f16_c_pytorch branch from 339b2d6 to a78ed36CompareJune 12, 2023 23:06
@cyanguwa

cyanguwa commented Jun 12, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

timmoon10
timmoon10 reviewed Jun 12, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/utils.cu Outdated
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/fused_attn.h
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
Comment threadtransformer_engine/pytorch/csrc/extensions.cu Outdated
@cyanguwa

cyanguwa commented Jun 12, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
update cudnn-frontend to v0.9.1
d3ceb5b
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Jun 13, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

zlsh80826
zlsh80826 reviewed Jun 13, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn.cpp Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
@timmoon10
timmoon10 self-requested a review June 13, 2023 18:35
@cyanguwa

cyanguwa commented Jun 14, 2023•
edited
Loading

Copy link
Copy Markdown
CollaboratorAuthor

@timmoon10@ksivaman@zlsh80826 Let me know if there's any other comments/suggestions. If all good, I think we can squash and merge. Thanks for reviewing!

zlsh80826
zlsh80826 reviewed Jun 20, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/cpp_extensions/fused_attn.py Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
Comment threadtransformer_engine/pytorch/cpp_extensions/fused_attn.py
Comment threadtransformer_engine/pytorch/cpp_extensions/fused_attn.py
@cyanguwa
cyanguwaforce-pushed the fused_attn/fp8_f16_c_pytorch branch 2 times, most recently from 71a1b95 to 0426febCompareJune 20, 2023 21:29
@cyanguwa
merge changes in main
d6344c2
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwacyanguwa reopened this Jun 20, 2023
@cyanguwa
change dropout from no to yes for fmha_v1
3c62660
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Jun 20, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
Merge branch 'main' into fused_attn/fp8_f16_c_pytorch
cab1ad6
@cyanguwa
remove output_rng_state in m512 kernels
b89c567
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa
fix elts_per_thread calculation in kvpacked fwd
6086419
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Jun 21, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@cyanguwa
remove dropout=0.0 restriction for m512 fused attn
868eb87
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Jun 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@zlsh80826
zlsh80826 self-requested a review June 22, 2023 05:35
zlsh80826
zlsh80826 reviewed Jun 22, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
zlsh80826
zlsh80826 reviewed Jun 22, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
zlsh80826
zlsh80826 reviewed Jun 22, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
zlsh80826
zlsh80826 approved these changes Jun 22, 2023
View reviewed changes
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
Comment threadtransformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu Outdated
@cyanguwa
remove output_rng_state completely from m512 kernels
8056018
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@cyanguwa

cyanguwa commented Jun 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@ksivaman

ksivaman commented Jun 22, 2023•
edited
Loading

Copy link
Copy Markdown
Member

Pipelines 8711280 and 8717530

ksivaman
ksivaman approved these changes Jun 22, 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.

LGTM :)

@ksivaman
ksivaman merged commit 5c58bea into NVIDIA:mainJun 22, 2023
ksivaman added a commit that referenced this pull request Jun 22, 2023
@cyanguwa@ksivaman
Add long sequence support for fused attention (#237)
92eabc3
* add long sequence support and unify three backends for fused attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* update cudnn-frontend to v0.9.1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* replace cpu_float2half_rn with __float2half_rn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix backend selection and NVTEDType
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fixes
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* make cudnn plan caches thread_local
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* replace cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix replacement of cuDNN throw with NVTE_CHECK
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* force dropout probablity to 0 in inference mode
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* change negInfinity to be consistent with m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove float2half conversion for scale_dropout
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* add back runtime api for sm detection
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* add gemm3 to enums FP8Fwd/BwdTensors
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* change dropout from no to yes for fmha_v1
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove output_rng_state in m512 kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix elts_per_thread calculation in kvpacked fwd
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove dropout=0.0 restriction for m512 fused attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove output_rng_state completely from m512 kernels
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwa
cyanguwa deleted the fused_attn/fp8_f16_c_pytorch branch February 22, 2024 00:34
@hungryGeek16hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
@ksivamanksivamanksivaman approved these changes
+1 more reviewer
@zlsh80826zlsh80826zlsh80826 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

0.10.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@timmoon10@zlsh80826