more style adjustments

ROCm · Jun 17, 2024 · 05c96a1 · 05c96a1
1 parent 0649fcd
commit 05c96a1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/csrc/custom/custom.cu b/csrc/custom/custom.cu
@@ -39,20 +39,18 @@ void LLMM1(at::Tensor in_a, at::Tensor in_b, at::Tensor out_c,
           at::cuda::getCurrentCUDAStream(), rows_per_block);
 }
 
-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);
+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);
 
-void wvSpltK(at::Tensor in_a, at::Tensor in_b, at::Tensor out_c,
-           const int N_in, const int CuCount) {
+void wvSpltK(at::Tensor in_a, at::Tensor in_b, at::Tensor out_c, const int N_in,
+             const int CuCount) {
   int M = in_a.size(0);
   int 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,
-          at::cuda::getCurrentCUDAStream(), CuCount);
+           at::cuda::getCurrentCUDAStream(), CuCount);
 }
 
-
-
 void LLGemmZZ(void* in_a, void* in_b, void* out_c, const int M, const int K,
               cudaStream_t stream, const int solidx);
 

diff --git a/vllm/model_executor/layers/tuned_gemm.py b/vllm/model_executor/layers/tuned_gemm.py
@@ -23,7 +23,8 @@ def __init__(self):
         self.bestsols = {}
         self.load_best_sols()
         self.create_ds()
-        self.CuCount = torch.cuda.get_device_properties(device='cuda').multi_processor_count
+        self.CuCount = torch.cuda.get_device_properties(
+            device='cuda').multi_processor_count
 
         if (self.save_gemm == 1):
             self.tuned_df = pd.DataFrame(columns=['M', 'N', 'K'])
@@ -89,7 +90,8 @@ def mm(self, inp, weights):
                 ]).drop_duplicates()
                 self.tuned_df.to_csv(self.untune_path, index=False)
 
-            if ((n == 4 or n == 3 or n== 2 or n == 1 ) and inp_view.dtype == torch.float16) :
+            if ((n == 4 or n == 3 or n== 2 or n == 1)
+                and inp_view.dtype == torch.float16):
                 out = torch.empty(inp_view.shape[0],
                               weights.shape[0],
                               dtype=inp_view.dtype,