Lint

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Author: jberchtold-nvidia

transformer_engine/jax/cpp_extensions/gemm.py

@@ -1623,21 +1623,20 @@ def lowering(
                 scaling_mode=scaling_mode.value,
                 is_grouped_dense_wgrad=is_grouped_dense_wgrad,
             )
-        else:
-            ffi_name = GroupedGemmPrimitive.name
-            return jax.ffi.ffi_lowering(ffi_name)(
-                ctx,
-                *args,
-                M=M,
-                N=N,
-                K=K,
-                lhs_is_trans=lhs_is_trans,
-                rhs_is_trans=rhs_is_trans,
-                scaling_mode=scaling_mode.value,
-                has_bias=has_bias,
-                is_grouped_dense_wgrad=is_grouped_dense_wgrad,
-                use_async_d2h_group_sizes=use_async_d2h_group_sizes,
-            )
+        ffi_name = GroupedGemmPrimitive.name
+        return jax.ffi.ffi_lowering(ffi_name)(
+            ctx,
+            *args,
+            M=M,
+            N=N,
+            K=K,
+            lhs_is_trans=lhs_is_trans,
+            rhs_is_trans=rhs_is_trans,
+            scaling_mode=scaling_mode.value,
+            has_bias=has_bias,
+            is_grouped_dense_wgrad=is_grouped_dense_wgrad,
+            use_async_d2h_group_sizes=use_async_d2h_group_sizes,
+        )
 
     @staticmethod
     def impl(

transformer_engine/jax/csrc/extensions/gemm.cpp

@@ -559,6 +559,7 @@ JAXX_GroupedTensorWrapper make_grouped_tensor(Buffer_Type const &data,
   return std::move(grouped_tensor_wrapper);
 }
 
+// This FFI is EXPERIMENTAL and subject to change without deprecation, intended for use in JAX's internal implementation of grouped GEMM.
 Error_Type GroupedGemmV2FFI(
     cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, Buffer_Type rhs_data,
     Buffer_Type rhs_sinv, Buffer_Type bias, Buffer_Type group_sizes, Buffer_Type alpha,