diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index dccecfc113..7181afa346 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -28,8 +28,18 @@ struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs ck_tile::index_t batch_stride_B_, ck_tile::index_t batch_stride_C_, ck_tile::index_t batch_count_) - : GemmHostArgs( - a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_, skip_a_lds_, skip_b_lds_), + : GemmHostArgs(a_ptr_, + b_ptr_, + c_ptr_, + k_batch_, + M_, + N_, + K_, + stride_A_, + stride_B_, + stride_C_, + skip_a_lds_, + skip_b_lds_), batch_stride_A(batch_stride_A_), batch_stride_B(batch_stride_B_), batch_stride_C(batch_stride_C_), @@ -103,7 +113,9 @@ struct BatchedGemmKernel : public GemmKernel; - using Traits = ck_tile::TileGemmTraits; - using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmUniversalTraits = + ck_tile::TileGemmUniversalTraits; using GemmPipelineProblem = ck_tile::GemmPipelineProblem; @@ -317,8 +323,19 @@ class TestCkTileGroupedGemm : public ::testing::Test // TODO add support for kbatch > 1 static constexpr ck_tile::index_t k_batch = 1; - gemm_descs.push_back( - {p_a, p_b, p_c, k_batch, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]}); + + gemm_descs.push_back({p_a, + p_b, + p_c, + k_batch, + M, + N, + K, + stride_As[i], + stride_Bs[i], + stride_Cs[i], + GroupedGemKernelParam::SkipALds, + GroupedGemKernelParam::SkipBLds}); } ck_tile::DeviceMem gemm_workspace;