[Feature]: Support pre-build multi arches for ck_gemm_moe_2stages_codegen

[yuguo68]

Suggestion Description

Currently ck_gemm_moe_2stages_codegen checks the gfx version at runtime,

aiter/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_ck2stages_common.py

Line 362 in ba9047e

arch = get_gfx()

.

Currently we don't use jit compile for aiter and prebuild all necessary kernels. It would be great to have the option to prebuild the moe_2stages kernels for both gfx942 and gfx950. There are also other kernels with the same issue.

Operating System

No response

GPU

MI300, MI350

ROCm Component

No response