Skinny gemm

CMakeLists.txt

@@ -677,6 +677,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
   # _rocm_C extension
   #
   set(VLLM_ROCM_EXT_SRC
+    "csrc/rocm/skinny_gemms.cu"
     "csrc/rocm/torch_bindings.cpp"
     "csrc/rocm/attention.cu")
 

csrc/rocm/ops.h

@@ -13,3 +13,9 @@ void paged_attention(torch::Tensor& out, torch::Tensor& exp_sums,
                      const std::optional<torch::Tensor>& alibi_slopes,
                      const std::string& kv_cache_dtype, torch::Tensor& k_scale,
                      torch::Tensor& v_scale);
+
+void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
+               at::Tensor& scale_a, at::Tensor& scale_b, const int64_t CuCount);
+
+torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
+                       const int64_t CuCount);