| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | #include <ATen/native/CPUBlas.h> |
| 5 | |
| 6 | // Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm |
| 7 | // If OpenBLS is available in the PyTorch wheel, we rely on it for fast |
| 8 | // bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path. |
| 9 | #if defined(VLLM_HAS_OPENBLAS) |
| 10 | extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k, |
| 11 | float* alpha, const at::BFloat16* a, int* lda, |
| 12 | const at::BFloat16* b, int* ldb, float* beta, float* c, |
| 13 | int* ldc); |
| 14 | |
| 15 | extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k, |
| 16 | float* alpha, const float* a, int* lda, const float* b, |
| 17 | int* ldb, float* beta, float* c, int* ldc); |
| 18 | |
| 19 | inline char blas_transpose(at::native::TransposeType trans) { |
| 20 | switch (trans) { |
| 21 | case at::native::TransposeType::NoTranspose: |
| 22 | return 'n'; |
| 23 | case at::native::TransposeType::Transpose: |
| 24 | return 't'; |
| 25 | case at::native::TransposeType::ConjTranspose: |
| 26 | return 'c'; |
| 27 | } |
| 28 | return 'n'; |
| 29 | } |
| 30 | |
| 31 | inline void blas_gemm(at::native::TransposeType transa, |
| 32 | at::native::TransposeType transb, int64_t m, int64_t n, |
| 33 | int64_t k, float alpha, const at::BFloat16* a, |
| 34 | int64_t lda, const at::BFloat16* b, int64_t ldb, |
| 35 | float beta, float* c, int64_t ldc) { |
| 36 | char transa_ = blas_transpose(transa); |
| 37 | char transb_ = blas_transpose(transb); |
| 38 | int m_ = static_cast<int>(m); |
| 39 | int n_ = static_cast<int>(n); |
| 40 | int k_ = static_cast<int>(k); |
| 41 | int lda_ = static_cast<int>(lda); |
| 42 | int ldb_ = static_cast<int>(ldb); |
| 43 | int ldc_ = static_cast<int>(ldc); |
| 44 | sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta, |
| 45 | c, &ldc_); |
| 46 | } |
| 47 | |
| 48 | inline void blas_gemm(at::native::TransposeType transa, |
| 49 | at::native::TransposeType transb, int64_t m, int64_t n, |
| 50 | int64_t k, float alpha, const float* a, int64_t lda, |
| 51 | const float* b, int64_t ldb, float beta, float* c, |
| 52 | int64_t ldc) { |
| 53 | char transa_ = blas_transpose(transa); |
| 54 | char transb_ = blas_transpose(transb); |
| 55 | int m_ = static_cast<int>(m); |
| 56 | int n_ = static_cast<int>(n); |
| 57 | int k_ = static_cast<int>(k); |
| 58 | int lda_ = static_cast<int>(lda); |
| 59 | int ldb_ = static_cast<int>(ldb); |
| 60 | int ldc_ = static_cast<int>(ldc); |
| 61 | sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta, |
| 62 | c, &ldc_); |
| 63 | } |
| 64 | |
| 65 | inline void blas_gemm(at::native::TransposeType, at::native::TransposeType, |
| 66 | int64_t, int64_t, int64_t, float, const at::Half*, |
| 67 | int64_t, const at::Half*, int64_t, float, float*, |
| 68 | int64_t) { |
| 69 | TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available."); |
| 70 | } |
| 71 | #else |
| 72 | template <typename scalar_t> |
| 73 | inline void blas_gemm(at::native::TransposeType transa, |
| 74 | at::native::TransposeType transb, int64_t m, int64_t n, |
| 75 | int64_t k, float alpha, const scalar_t* a, int64_t lda, |
| 76 | const scalar_t* b, int64_t ldb, float beta, float* c, |
| 77 | int64_t ldc) { |
| 78 | auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT; |
| 79 | gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k, |
| 80 | at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc); |
| 81 | } |
| 82 | #endif |