diff --git a/.github/container/test-maxtext.sh b/.github/container/test-maxtext.sh
index 0377237f5..0dc26c8c1 100755
--- a/.github/container/test-maxtext.sh
+++ b/.github/container/test-maxtext.sh
@@ -232,8 +232,7 @@ export BASE_XLA_FLAGS=${BASE_XLA_FLAGS:---xla_gpu_enable_latency_hiding_schedule
                 --xla_gpu_enable_pipelined_all_gather=true 
                 --xla_gpu_enable_pipelined_reduce_scatter=true 
                 --xla_gpu_enable_pipelined_all_reduce=true 
-                --xla_gpu_enable_while_loop_double_buffering=true
-                --xla_gpu_enable_triton_softmax_fusion=false 
+                --xla_gpu_enable_while_loop_double_buffering=true 
                 --xla_gpu_enable_all_gather_combine_by_dim=false 
                 --xla_gpu_enable_reduce_scatter_combine_by_dim=false 
                 --xla_disable_hlo_passes=rematerialization}
diff --git a/rosetta/docs/PGLE.md b/rosetta/docs/PGLE.md
index dcf925fb9..204fb4efa 100644
--- a/rosetta/docs/PGLE.md
+++ b/rosetta/docs/PGLE.md
@@ -70,7 +70,6 @@ export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
 --xla_gpu_enable_pipelined_reduce_scatter=true
 --xla_gpu_enable_pipelined_all_reduce=true
 --xla_gpu_enable_while_loop_double_buffering=true
---xla_gpu_enable_triton_softmax_fusion=false
 --xla_gpu_enable_all_gather_combine_by_dim=false
 --xla_gpu_enable_reduce_scatter_combine_by_dim=false
 --xla_disable_hlo_passes=rematerialization
diff --git a/rosetta/rosetta/projects/maxtext/README.md b/rosetta/rosetta/projects/maxtext/README.md
index b137edfd0..de2d60893 100644
--- a/rosetta/rosetta/projects/maxtext/README.md
+++ b/rosetta/rosetta/projects/maxtext/README.md
@@ -77,7 +77,6 @@ XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
             --xla_gpu_enable_pipelined_reduce_scatter=true 
             --xla_gpu_enable_pipelined_all_reduce=true 
             --xla_gpu_enable_while_loop_double_buffering=true
-            --xla_gpu_enable_triton_softmax_fusion=false 
             --xla_gpu_enable_all_gather_combine_by_dim=false 
             --xla_gpu_enable_reduce_scatter_combine_by_dim=false 
             --xla_disable_hlo_passes=rematerialization"
diff --git a/rosetta/rosetta/projects/maxtext/scripts/example_slurm.sub b/rosetta/rosetta/projects/maxtext/scripts/example_slurm.sub
index 93894c75d..fd45555b7 100644
--- a/rosetta/rosetta/projects/maxtext/scripts/example_slurm.sub
+++ b/rosetta/rosetta/projects/maxtext/scripts/example_slurm.sub
@@ -62,7 +62,6 @@ export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true
                 --xla_gpu_enable_pipelined_reduce_scatter=true
                 --xla_gpu_enable_pipelined_all_reduce=true
                 --xla_gpu_enable_while_loop_double_buffering=true
-                --xla_gpu_enable_triton_softmax_fusion=false
                 --xla_gpu_enable_all_gather_combine_by_dim=false
                 --xla_gpu_enable_reduce_scatter_combine_by_dim=false
                 --xla_disable_hlo_passes=rematerialization"