diff --git a/transformer_vm/model/transformer.cpp b/transformer_vm/model/transformer.cpp index b287f61..9ba0f10 100644 --- a/transformer_vm/model/transformer.cpp +++ b/transformer_vm/model/transformer.cpp @@ -22,6 +22,8 @@ #ifdef __APPLE__ #define ACCELERATE_NEW_LAPACK #include +#elif defined(USE_OPENBLAS) +#include #endif #include "hull2d_cht.h" @@ -158,7 +160,7 @@ static inline void add_position_encoding(double* x, int pos) { static inline void matvec(const double* __restrict__ W, const double* __restrict__ x, double* __restrict__ y, int rows, int cols) { -#if defined(__APPLE__) && !defined(NO_BLAS) +#if (defined(__APPLE__) || defined(USE_OPENBLAS)) && !defined(NO_BLAS) cblas_dgemv(CblasRowMajor, CblasNoTrans, rows, cols, 1.0, W, cols, x, 1, 0.0, y, 1); #else diff --git a/transformer_vm/runner.py b/transformer_vm/runner.py index 6a6b84e..1ae186b 100644 --- a/transformer_vm/runner.py +++ b/transformer_vm/runner.py @@ -113,6 +113,16 @@ def _build_cpp_engine(): ] else: cmd = ["g++", "-std=c++17", "-O3", "-I", attn_dir, source, "-o", binary] + # Opt into OpenBLAS for the matvec path if cblas.h is discoverable. + # Falls back silently to the scalar nested loop when not installed. + try: + pkg = subprocess.run( + ["pkg-config", "--cflags", "--libs", "openblas"], + capture_output=True, text=True, check=True, + ) + cmd += ["-DUSE_OPENBLAS"] + pkg.stdout.split() + except (subprocess.CalledProcessError, FileNotFoundError): + pass try: subprocess.check_call(cmd) logger.info("[engine] Built: %s", binary)