update

Author: zhiding512

3rdparty/composable_kernel

@@ -1,10 +1,11 @@
 /* SPDX-License-Identifier: MIT
    Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
 */
-#include "rocm_ops.hpp"
 #include "moe_op.h"
+#include "rocm_ops.hpp"
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
+      AITER_ENUM_PYBIND;
       MOE_OP_PYBIND;
 }

csrc/pybind/moe_op_pybind.cu

@@ -668,8 +668,8 @@ def weight_per_128x128_quant(weight, quant_dtype):
 torch.cuda.manual_seed_all(seed)
 l_dtype = ["bf16", "fp16"][:1]
 # l_dim = [(6144, 4096)]
-l_dim = [(7168, 256)]
-# l_dim = [(3072, 3072)]
+# l_dim = [(7168, 256)]
+l_dim = [(3072, 3072)]
 l_tokenNum = [
     # 1,
     # 2,
@@ -693,8 +693,8 @@ def weight_per_128x128_quant(weight, quant_dtype):
     # (aiter.QuantType.per_Token, dtypes.fp8, dtypes.fp8),  # a8w8
     # (aiter.QuantType.per_Token, dtypes.fp8, torch.int4),  # a8w4
     # (aiter.QuantType.per_1x32, dtypes.fp4x2, dtypes.fp4x2),  # a4w4
-    (aiter.QuantType.per_128x128, dtypes.fp8, dtypes.fp8),  # a8w8
-    # (aiter.QuantType.per_1x32, dtypes.bf16, dtypes.fp4x2),  # a16w4
+    # (aiter.QuantType.per_128x128, dtypes.fp8, dtypes.fp8),  # a8w8
+    (aiter.QuantType.per_1x32, dtypes.bf16, dtypes.fp4x2),  # a16w4
 ]
 l_act = [aiter.ActivationType.Silu, aiter.ActivationType.Gelu][:1]
 l_doweight_stage1 = [False, True][:1]