[TE/JAX] Disable FusedAttn with FFI by default (#1298)

* disable fused attn with ffi --------- Signed-off-by: Phuong Nguyen <[email protected]>
NVIDIA · Oct 31, 2024 · 23caab3 · 23caab3
1 parent 9dddb36
commit 23caab3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -379,7 +379,7 @@ def lowering(
 
         wkspace_aval = ctx.avals_out[-1]
 
-        if is_ffi_enabled():
+        if is_ffi_enabled() and bool(os.getenv("NVTE_JAX_FUSED_ATTN_WITH_FFI", "0")):
             name = "te_fused_attn_forward_ffi"
             out = ffi.ffi_lowering(name)(
                 ctx,