(tvm) 19/19: fix

[masahi]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit c796f47f95e18b838281940a138200a711b2f741
Author: Wuwei Lin <wu...@apache.org>
AuthorDate: Tue Dec 12 17:51:57 2023 +0000

    fix
---
 src/runtime/contrib/cutlass/weight_preprocess.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/runtime/contrib/cutlass/weight_preprocess.cc 
b/src/runtime/contrib/cutlass/weight_preprocess.cc
index f82591fa84..a1508756b0 100644
--- a/src/runtime/contrib/cutlass/weight_preprocess.cc
+++ b/src/runtime/contrib/cutlass/weight_preprocess.cc
@@ -39,7 +39,7 @@ namespace runtime {
 TVM_REGISTER_GLOBAL("cutlass.ft_preprocess_weight")
     .set_body_typed([](NDArray packed_weight, int sm, bool is_int4) {
       bool is_2d = packed_weight->ndim == 2;
-      int num_experts = 1;
+      int num_experts = is_2d ? 1 : packed_weight->shape[0];
       int rows = packed_weight->shape[is_2d ? 0 : 1];
       int cols = packed_weight->shape[is_2d ? 1 : 2];