Add KBatch

ROCm · Oct 9, 2024 · 24a3079 · 24a3079
1 parent 15c61bb
commit 24a3079
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/aten/src/ATen/native/hip/ck_gemm_template.h b/aten/src/ATen/native/hip/ck_gemm_template.h
@@ -170,6 +170,8 @@ void gemm_impl(CUDABLAS_GEMM_ARGTYPES(Dtype)) {
   int StrideB = ldb;
   int StrideC = ldc;
 
+  int KBatch = 1;
+
   float falpha = alpha;
   float fbeta = beta;
 
@@ -268,6 +270,7 @@ void gemm_impl(CUDABLAS_GEMM_ARGTYPES(Dtype)) {
      StrideA,
      std::array<ck::index_t, 0>{},
      StrideC,
+     KBatch,
      a_element_op,
      b_element_op,
      c_element_op);