From 05c96a1c669d0bc2126c2a50c08b0c3b38f68d21 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 17 Jun 2024 21:40:18 +0000 Subject: [PATCH] more style adjustments --- csrc/custom/custom.cu | 12 +++++------- vllm/model_executor/layers/tuned_gemm.py | 6 ++++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/csrc/custom/custom.cu b/csrc/custom/custom.cu index 5c8beed37b304..9e92187967d47 100644 --- 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 index 5e1f4f35d3d87..8b9b20722fb9e 100644 --- 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,