buildkit option for docker build

(cherry picked from commit f075f90)

Author: kiran-thumma

build/rocm/ci_build

@@ -54,6 +54,8 @@ def dist_wheels(
     except subprocess.CalledProcessError as e:
         print(f"Failed to remove Docker image {image}: {e}")
 
+    os.environ["DOCKER_BUILDKIT"] = "1"
+
     cmd = [
         "docker",
         "build",
@@ -166,6 +168,8 @@ def dist_docker(
 
     md = _fetch_jax_metadata(xla_path)
 
+    os.environ["DOCKER_BUILDKIT"] = "1"
+
     cmd = [
         "docker",
         "build",