From 080b6084281ca7a08dd61990c14060bfc3d0b41a Mon Sep 17 00:00:00 2001
From: ashors1 <ashors@nvidia.com>
Date: Tue, 23 Apr 2024 13:52:27 -0700
Subject: [PATCH] enable TE fused attention by default in rosetta pax container

---
 rosetta/Dockerfile.pax | 1 +
 1 file changed, 1 insertion(+)

diff --git a/rosetta/Dockerfile.pax b/rosetta/Dockerfile.pax
index 140ab2b5f..c702aa530 100644
--- 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})