fix assertion bug for SWA API in TE-JAX (#1242)

fixed assertion bug for SWA Signed-off-by: Md Fahim Faysal Khan <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]> Co-authored-by: Phuong Nguyen <[email protected]>
NVIDIA · Oct 16, 2024 · 43b9e1e · 43b9e1e
1 parent f6b766b
commit 43b9e1e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -1061,7 +1061,7 @@ def partition(config, mesh, arg_infos, result_infos):
         # Call base implementation for non-context parallel mesh to avoid unecessary work.
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         assert (
-            is_context_parallel and config.window_size[0] == -1
+            not is_context_parallel or config.window_size[0] == -1
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
@@ -1158,7 +1158,7 @@ def partition(config, mesh, arg_infos, result_infos):
         # Call base implementation for non-context parallel mesh to avoid unecessary work.
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         assert (
-            is_context_parallel and config.window_size[0] == -1
+            not is_context_parallel or config.window_size[0] == -1
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnBwdPrimitive.partition(config, mesh, arg_infos, result_infos)