From 52cc27d1352923f7aac7acc157191d2b1fe30980 Mon Sep 17 00:00:00 2001
From: Wuwei Lin <wuwei@apache.org>
Date: Fri, 8 Mar 2024 18:22:54 -0800
Subject: [PATCH] [Cutlass] Fix shape in fp8 group gemm (#58)

---
 src/runtime/contrib/cutlass/fp8_group_gemm.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/runtime/contrib/cutlass/fp8_group_gemm.cu b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
index c93da6ff5766..4cc6192f8e13 100644
--- a/src/runtime/contrib/cutlass/fp8_group_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
@@ -56,7 +56,7 @@ void tvm_cutlass_fp8_group_gemm(NDArray x, NDArray weight, NDArray indptr, NDArr
   CHECK_EQ(alpha->dtype.bits, 32);
   int num_groups = weight->shape[0];
   int n = weight->shape[1];
-  int k = weight->shape[2];
+  int k = x->shape[1];
   const float* beta = nullptr;
   cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
   cutlass_group_gemm(static_cast<ElementA*>(x->data), static_cast<ElementB*>(weight->data),