enable TE fused attention by default in rosetta pax container

NVIDIA · Apr 23, 2024 · 080b608 · 080b608
1 parent fbb9634
commit 080b608
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rosetta/Dockerfile.pax b/rosetta/Dockerfile.pax
@@ -24,6 +24,7 @@ ARG UPDATE_PATCHES
 ARG UPDATED_TE_REF
 
 ENV ENABLE_TE=1
+ENV NVTE_FUSED_ATTN=1
 
 RUN --mount=target=/mnt/jax-toolbox,from=jax-toolbox <<"EOF" bash -exu
 MANIFEST_DIR=$(dirname ${MANIFEST_FILE})