') + ')', '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); } })(); })(); Flash attn gqa by lhb8125 · Pull Request #365 · NVIDIA/TransformerEngine · GitHub
Skip to content

Flash attn gqa - #365

Closed
lhb8125 wants to merge 56 commits into
NVIDIA:mainfrom
lhb8125:flash_attn_GQA
Closed

Flash attn gqa#365
lhb8125 wants to merge 56 commits into
NVIDIA:mainfrom
lhb8125:flash_attn_GQA

Conversation

@lhb8125

@lhb8125lhb8125 commented Aug 8, 2023

Copy link
Copy Markdown
Contributor

Fix bug when using GQA, where the shape of query layer is different from key/value layer.

ptrendxand others added 30 commits October 20, 2022 11:42
@ptrendx
Change version to 0.2.0
29e9e27
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@ptrendx
Changing version to 0.3
c15f447
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@ksivaman@ptrendx
Full activation recompute checkpointing bug fix (NVIDIA#31)
73166c4
fix checkpoint loading bug for FAR
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ptrendx
Change version to 0.4
c2f15d5
@ptrendx
Change version to 0.5
7d10144
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@ksivaman@ptrendx
Address steady memory increase and bloated checkpoints (NVIDIA#63)
126232d
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Change version to 0.6
04cc731
@ksivaman@ptrendx
3rd party acknowledgements (NVIDIA#82)
ce58fc2
add 3rd party acknowledgements
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@ptrendx
Fix unfused QKV params case; stack vs interleave option (NVIDIA#83)
4c35891
* fix qkv weight unfused path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix non FA non interleaved case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@ptrendx
Fix flash attention (NVIDIA#84)
bb12038
* ignore self attention mask for causal type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* further relax checks to run FA, update docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix pytorch softmax path
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* minimum ampere requirement for fa
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@ptrendx
fix bug in non-FP8 nvfuser path (NVIDIA#81)
f18e677
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ptrendx
Change version to 0.7
2bc7dfe
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@ptrendx
Add SECURITY.md (NVIDIA#110)
f4955d3
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@ksivaman@ptrendx
Catch FA internal error with compute capability 8.6 (NVIDIA#113)
e5ab211
FA doesn't support compute 8.6 with head_dim>64
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@ptrendx
Fix usage of return_bias argument (NVIDIA#114)
7e8c3e6
* fix usage of return_bias argument
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@timmoon10Naman Goyal
@ksivaman@ptrendx
Fix zombie process when querying TE install path (NVIDIA#121)
626da0d
* Remove zombie process from querying TE install path
Co-authored-by: Naman Goyal <naman@fb.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix FA version checking
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix unused import error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix lint warning
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Naman Goyal <naman@fb.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@jeng1220@ptrendx
[JAX] Add TE examples (NVIDIA#108)
084b1e5
* refactor JAX examples
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix doc-string
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add dp example
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* refactor
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix params_axes_pspec
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Add model parallel example and refactor
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* align code and readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update verification
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add mask
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* num_gpu is configurable
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* solvepylint issue
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* ignore markdown and txt file from license check
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Update README.md
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add flax into requirements.txt
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@mingxu1067@ptrendx
Fix Bugs of TE/JAX (NVIDIA#119)
bc0e448
* Support transpose_bs when decoded=True
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix Bugs,
1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix wrong masks in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed wrong assert condition in TransformerLayer
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix amax is not set as 0 in each step.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Enhance rules conflict checking and docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* fix code formatting.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
@ksivaman@nzmora-nvidia@ptrendx
Change FP8 recipe defaults (NVIDIA#112)
78c375d
* Change FP8 recipe defaults
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Increase default amax history length
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Always check history size
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* no amax history for onnx export
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* revert onnx export test changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix indices in onnx test
Co-authored-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Neta Zmora <nzmora@nvidia.com>
@ksivaman@ptrendx
Add FP8 support for Ada (NVIDIA#129)
a753715
* Add FP8 support for Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* better message
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Address review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* better message for no fp8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* same thing for onnx test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix CI and review
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@nouiz@ksivaman@ptrendx
Update installation instruction for JAX and add some dependencies. (N…
770e968
…VIDIA#117)
* Update installation instructio for JAX and add some depenencies.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Bring back support for none pip installed pybind11.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Apply suggestions from code review
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>
* Changes following review.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Change order to make it more clear.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Add other reviers suggestion.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* pybind11 is needed for all FW.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Add flax as a dep
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
* Update README.rst
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>
---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ptrendx
Change version to 0.8
414fd73
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@erhoo82@ksivaman@ptrendx
Amax reduction interval (NVIDIA#154)
ee87982
* amax reduction internval
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Skip TP-domain only AMAX reduction when TP-group is not initialized
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* Update transformer_engine/pytorch/fp8.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* check TP group initialized
Signed-off-by: Sangkug Lym <slym@nvidia.com>
fix
Signed-off-by: Sangkug Lym <slym@nvidia.com>
---------
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@erhoo82
@vasunvidia@ptrendx
TP communication overlap with userbuffers (NVIDIA#147)
e64fc3b
* Port initial changes
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* readd FA include for PyTorch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Re-enable sm_70 + cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* LICENSE, cleanup header
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* 5k -> 173 errors
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* license and fixes in userbuffers-host
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* next round fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* final cpp cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* pylinting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix from linting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Turn off default async amax reduction (NVIDIA#148)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* remove unused code path
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* cleanup Macros
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* fix conflict resolution bug
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* Fix gencode flags in setup (NVIDIA#145)
* Fix gencode flags based on cuda version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* review suggestions
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* revert append_nvcc_threads change
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Change overlap config dict error message
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* simplify ub initialization
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix sanity imports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cpplint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix TensorFlow build
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix TE macros in public header
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* More fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* compiles with and w/o MPI
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fixes for python side annotations for conditional compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* link gdrAPI only when MPI found
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix comments for dummy var
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix linking
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* load MPI before TE
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add Py side argument checks
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* remove unused code and catch silent failures
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix cpp tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix find_lib path for tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
@ksivaman@erhoo82
@timmoon10@ptrendx
Move userbuffer to PyTorch (NVIDIA#162)
7bf886d
* Initial refactor; linker error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix linking issue and make mpi conditional
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix TF/JAX build
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Use max SMs at the last RS chunk in pipelined overlap
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Make userbuffers support opt-in
Decouple userbuffers from MPI. Refactor MPI handling in build system. Standardize names to "userbuffers".
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Lint
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
@cyanguwa@ksivaman@ptrendx
Add FP8 fused attention (NVIDIA#155)
ac5d44e
* Add FP8 fused attention to TE for PyTorch
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add license for cudnn-frontend, modify installation requirements, and refactor some headers for aesthetics
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add c api docs for fused attention
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add exception for unsupported precision/sequence length combinations
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix installation requirement for non fused attn use cases
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix docs for fused-attn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* prefix enums with NVTE_ and replace old MHA_Matrix with NVTE_QKV_Matrix
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes based on PR comments
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description for kvpacked fwd
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix description of Bias in C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* minor fixes for cudnn requirement and description for QKV tensors
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix QKV layout description and support matrix for C api
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* add asserts to cpp_extensions for qkv layout/bias type/attn mask type
Signed-off-by: Charlene Yang <charleney@nvidia.com>
* fix typo precision
Signed-off-by: Charlene Yang <charleney@nvidia.com>
---------
Signed-off-by: Charlene Yang <charleney@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Charlene Yang <charleney@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@erhoo82@ksivaman@ptrendx
zero inter-node communication buffer (NVIDIA#163)
e1ef756
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman@ptrendx
Remove userbuf docs (NVIDIA#164)
9d90eb4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Faster split of QKV for FlashAttention (NVIDIA#166)
71488db
* Faster split of QKV for FlashAttention
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* CI fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Fix
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Message with assert
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Review comments
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* review
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix misalignment error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* make clarifying comment and check strides
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Remove the nonexistent parameter from fused attention documentation (N…
87706dc
…VIDIA#181)
* Remove the nonexistent parameter from fused attention documentation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Remove the second instance
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
mingxu1067and others added 25 commits May 2, 2023 09:44
@mingxu1067@ksivaman@ptrendx
[JAX] Adjust Module Structure. (NVIDIA#169)
00707bb
* Adjust Module Structure.
1. Collect Flax related modules to a sub-folder, flax.
2. Add a function to unify scale_init for zero-centered-gamma LN.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Make changes be compatible to previous versions.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adapt jax/examples to the new module structure.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Update jax/docs and Add deprecated warning.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Update README
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding deprecated_wrapper
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Adding deprecated warning to flax modules which imported via transformer_engine.jax
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix CI errors and update docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Removing unnecessary deprecated warning in docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Implementing __iter__ to DeprecatedEnum.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwa@ptrendx
Correct cuDNN version requirement (NVIDIA#184)
550da28
correct cuDNN version requirement
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@erhoo82@ksivaman@ptrendx
Use separate streams for pushsend/recv kernels in UB p2p exchanges (N…
d3d4191
…VIDIA#188)
* using different strems for pushsend and pushrecv
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* fix stream dependency
Signed-off-by: Sangkug Lym <slym@nvidia.com>
* add wait from main_stream to memcpy stream
Signed-off-by: Sangkug Lym <slym@nvidia.com>
---------
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ShriyaRishab@erhoo82
Shriya/tp overlap patch (NVIDIA#205)
8e5f00f
userbuffer pushsend/recv fix with atomicAdd_system
Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Sangkug Lym <slym@nvidia.com>
@ptrendx
Change VERSION to 0.9.0
2e81922
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
@cyanguwa@ptrendx
Relax checks for attn_mask_type in FlashAttention (NVIDIA#226)
f92c430
* relax attn mask type checks for FlashAttention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* disable flash attn if mask tensor is not None
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix the logic for flash attn
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* minor fix for lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
@zlsh80826@ptrendx
Jax bug fixes for the dot product attention (NVIDIA#236)
06cacd2
* Unfused scale+softmax if bias is present
Signed-off-by: Reese Wang <rewang@nvidia.com>
* WAR a causal masking + no_bias bug and add the unittests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix the optional args (bias) sharding
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Disable fused attn in JAX by default, enable it with NVTE_USE_FUSED_ATTN
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add thread local for the plan cache
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename dbeta to dbias for the readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add scaled softmax with dropout test cases
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Updated NVTE_FUSED_ATTN variable name
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
@ksivaman
Change VERSION to 0.10.0
5a594cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@zlsh80826@ksivaman
[JAX] Add self_attn_mask_type and replace attn_type (NVIDIA#273)
84a4a75
* Add self_attn_mask_type and replace attn_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the keyword style for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace attn_type with attn_mask_type in praxis transformer
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix typos
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@zlsh80826@ksivaman
Support dropout for the fused attention when max seqlen <= 512 (NVIDI…
4244ba9
…A#227)
* Enable fused attention dropout
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Cast the uint32 key/counter to int64
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update dropout support in fused attention docs
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revise devPtrCuSeqlen* to align the naming
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Support different Jax PRNG impls
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert CastAsync since it is not used
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Implement is_training for 16-bit fused attn
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Add fused attn with dropout sanity unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the comments readability and rng_state checker
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the attention dropout shape to align other frameworks
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Make encoder tests deterministic
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Change the default seed for the jax encoder tests
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Maintain offset in TE
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Enhance the resource safety
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Revert rng_state type to allow only i64
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Handle the corner case for elts_per_threads calculation
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Populate rng state by kernels
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Rename rng_state as seed in cpp_extensions
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update the attention dropout comment
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwa@ksivaman
Add long sequence support for fused attention (NVIDIA#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>
@asfiyab-nvidia@ksivaman
Fix BF16 ONNX export for successful ONNX Runtime Verification (NVIDIA…
ac919e4
…#290)
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
@nzmora-nvidia@ksivaman
Fix layer_norm ONNX export (NVIDIA#293)
96ed6fc
* Fix ONNX export of layer_norm
ONNX has a spec bug: ConstantOfShape supports all dtypes except for BF16.
To WAR we use dtype FP32 and then cast to BF16.
Will also issue a PR to the ONNX sig committee to change the spec in opset 20.
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
* fix lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Neta Zmora <nzmora@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Changed the version to 0.11.0
cc0831d
@zlsh80826@ksivaman
[JAX] Fully remove attn_type and set self_attn_mask_type default to '…
94beb13
…causal' (NVIDIA#324)
* Fully remove attn_type and set self_attn_mask_type default to 'causal'
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix tests with new arguments
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Explicit self_attn_mask_type for examples
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>
* Update transformer_engine/jax/flax/transformer.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: zlsh80826 <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
FA does not support head_dim > 64 on Ada (NVIDIA#328)
32ad922
* FA does not support head_dim > 64 on Ada
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add cc8.7 to no FA list
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
FlashAttention 2.0 support (NVIDIA#329)
33576be
* FA v2.0 support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix typo
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Relax FA 2.0 checks for Ada (NVIDIA#331)
0777408
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Make QK layer scaling opt-in (NVIDIA#339)
3f9db84
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ptrendx@ksivaman
Exposing RMSNorm in pyTorch (NVIDIA#306)
058f912
* Exposing RMSNorm in pyTorch extensions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* First pass at the Python API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Small fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Added numerics tests and fixed issues
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Lint fixes
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Added RMSNorm to LayerNormMLP
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Added ONNX export and tests for RMSNorm
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Fix python lint
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Fix BERT case
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Added normalization option to the TransformerLayer
Added tests
Fixed test failures
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
* Fix documentation
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix kwarg bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix IMA and invalid type error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Increase RMSNorm threshold for bf16 case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@cyanguwa@ksivaman
Add support for multi-query and grouped-query attention (NVIDIA#338)
5ed7e82
* add support for multi-query/grouped-query attention
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix lint
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* revert to flash-attn 1.0.6 and build 2.0.0.post1 manually in CI
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* add keyword name for DPA input
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix fused attn tests
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* fix skipif for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* Update transformer_engine/pytorch/attention.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Update tests/pytorch/test_fused_attn.py
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix TP and SP case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* add skipifs for pytest
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* remove higher limit for flash-attn version
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>
@ksivaman
Add compilation OOM note for FA 2.0 (NVIDIA#346)
9347b10
Add compilation warning for FA 2.0
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@timmoon10@ksivaman
Replace deprecated sharding API in JAX test (NVIDIA#332)
3f01b4f
Replace deprecated sharding API
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@lhb8125
fix flash_attn interface when using GQA
6322a4d
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>
@lhb8125
Merge branch 'main' into flash_attn_GQA
e3dfec5
@ksivaman

ksivaman commented Aug 8, 2023

Copy link
Copy Markdown
Member

@lhb8125 Could you rebase your changes with main and also sign-off your commit?

@lhb8125

lhb8125 commented Aug 9, 2023

Copy link
Copy Markdown
ContributorAuthor

@lhb8125 Could you rebase your changes with main and also sign-off your commit?

Hi, since the history of my branch is a little messy, I create another clean PR, please review that.

@lhb8125

lhb8125 commented Aug 9, 2023

Copy link
Copy Markdown
ContributorAuthor

@ksivaman

@ksivaman

ksivaman commented Aug 9, 2023

Copy link
Copy Markdown
Member

Closing in favor of #367

@ksivamanksivaman closed this Aug 9, 2023
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

13 participants

@lhb8125@ksivaman@ptrendx@timmoon10@jeng1220@mingxu1067@nouiz@erhoo82@cyanguwa@ShriyaRishab@zlsh80826@asfiyab-nvidia@nzmora-nvidia