minor clean up in setup

setup.py

@@ -139,7 +139,8 @@ def validate_and_update_archs(archs):
 # We want this even if SKIP_CUDA_BUILD because when we run python setup.py sdist we want the .hpp
 # files included in the source distribution, in case the user compiles from source.
 if IS_ROCM:
-    subprocess.run(["git", "submodule", "update", "--init", "csrc/composable_kernel"])
+    if not USE_TRITON_ROCM:
+        subprocess.run(["git", "submodule", "update", "--init", "csrc/composable_kernel"])
 else:
     subprocess.run(["git", "submodule", "update", "--init", "csrc/cutlass"])
 
@@ -304,17 +305,6 @@ def validate_and_update_archs(archs):
         )
     )
 elif not SKIP_CUDA_BUILD and IS_ROCM:
-    ck_dir = "csrc/composable_kernel"
-
-    #use codegen get code dispatch
-    if not os.path.exists("./build"):
-        os.makedirs("build")
-
-    os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd --output_dir build --receipt 2")
-    os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_appendkv --output_dir build --receipt 2")
-    os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv --output_dir build --receipt 2")
-    os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d bwd --output_dir build --receipt 2")
-
     print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
     TORCH_MAJOR = int(torch.__version__.split(".")[0])
     TORCH_MINOR = int(torch.__version__.split(".")[1])
@@ -323,6 +313,17 @@ def validate_and_update_archs(archs):
         # Skip C++ extension compilation if using Triton Backend
         pass
     else:
+        ck_dir = "csrc/composable_kernel"
+
+        #use codegen get code dispatch
+        if not os.path.exists("./build"):
+            os.makedirs("build")
+
+        os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd --output_dir build --receipt 2")
+        os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_appendkv --output_dir build --receipt 2")
+        os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv --output_dir build --receipt 2")
+        os.system(f"{sys.executable} {ck_dir}/example/ck_tile/01_fmha/generate.py -d bwd --output_dir build --receipt 2")
+
         # Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
         # See https://github.com/pytorch/pytorch/pull/70650
         generator_flag = []