Add praxis grok DPA patch (#889)

NVIDIA · Jun 11, 2024 · bcb7205 · bcb7205
1 parent 98a30fe
commit bcb7205
Show file tree

Hide file tree

Showing 9 changed files with 775 additions and 179 deletions.
diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml
@@ -48,6 +48,7 @@ praxis:
   patches:
     pull/27/head: file://patches/praxis/PR-27.patch # This PR allows XLA:GPU to detect the MHA pattern more easily to call fused kernels from cublas.
     pull/36/head: file://patches/praxis/PR-36.patch # adds Transformer Engine support
+    pull/74/head: file://patches/praxis/PR-74.patch # experimental support for using TE FMHA in GQA
 lingvo:
   # Used only in ARM pax builds
   url: https://github.com/tensorflow/lingvo.git
@@ -179,4 +180,4 @@ orbax-checkpoint:
   url: https://github.com/google/orbax.git
   tracking_ref: main
   latest_verified_commit: 16c2d409e365576284dbaf190ac002b24c1f927f
-  mode: pip-vcs
+  mode: pip-vcs