From 82bb8dde6e8e2f30579adc795248003482601fcf Mon Sep 17 00:00:00 2001
From: Jing Zhang <jizhan@fb.com>
Date: Mon, 21 Oct 2024 12:42:15 -0700
Subject: [PATCH] fixed splitk

---
 .../gpu/element/unary_element_wise_operation.hpp              | 1 -
 .../gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp                | 4 ++--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
index a6f75279fb..a0dc29e837 100644
--- a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
+++ b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
@@ -59,7 +59,6 @@ __device__ inline half2_t pki4_to_half2(pk_i4_t q)
     int x_h = (x_u8 & 0xf0) << 12;
 
     const int EX = 0x64006400;
-
     const int SUB = 0xE408E408; //-8
 
     int lo = (x_l | x_h) | EX;
diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
index 4a7695ed1d..90b089c4cf 100644
--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp
@@ -601,7 +601,7 @@ struct GridwiseGemm_xdl_cshuffle_v3
         {
             if constexpr(is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
             {
-                a_k_split_offset = blockIdx.z * karg.KRead;
+                a_k_split_offset = blockIdx.z * karg.KRead / APackedSize;
             }
             else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)
             {
@@ -614,7 +614,7 @@ struct GridwiseGemm_xdl_cshuffle_v3
             }
             else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, BLayout>)
             {
-                b_k_split_offset = blockIdx.z * karg.KRead;
+                b_k_split_offset = blockIdx.z * karg.KRead / BPackedSize;
             }
 
             if(blockIdx.z < static_cast<uint32_t>(karg.KBatch - 1))