BF16 Skinny Optimization

csrc/rocm/custom.cu

@@ -37,14 +37,16 @@ void LLMM1(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
 }
 
 void wvSpltK_(void* in_a, void* in_b, void* out_c, const int M, const int K,
-              const int N, cudaStream_t stream, const int CuCount);
+              const int N, const int Itp_in, cudaStream_t stream,
+              const int CuCount);
 
 void wvSpltK(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
-             const int64_t N_in, const int64_t CuCount) {
+             const int64_t N_in, const int64_t Itp_in, const int64_t CuCount) {
   auto M = in_a.size(0);
   auto K = in_a.size(1);
   int N = N_in;
-  wvSpltK_(in_a.data_ptr(), in_b.data_ptr(), out_c.data_ptr(), M, K, N,
+  int Itp = Itp_in;
+  wvSpltK_(in_a.data_ptr(), in_b.data_ptr(), out_c.data_ptr(), M, K, N, Itp,
            at::cuda::getCurrentCUDAStream(), CuCount);
 }