Commit a7b94df2c for llama.cpp
commit a7b94df2c616bc1f62a73b964b4a71cb0dcc488e
Author: SXX <song_xiaoxi@126.com>
Date: Mon Oct 5 03:22:17 2026 +0800
ggml-cpu: support BF16/FP16/FP32 K tails in tinyBLAS on x86 (#29806)
* ggml-cpu: vectorize BF16 K tails in tinyBLAS
* tests: Skip tinyBLAS when use_ref is enabled so CPU tests compare against the vec_dot path.
* ggml-cpu: vectorize tinyBLAS F16/F32 tails
diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
index 8620c6c7a..7c274c03b 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.c
+++ b/ggml/src/ggml-cpu/ggml-cpu.c
@@ -1305,7 +1305,7 @@ void ggml_compute_forward_mul_mat(
const bool src1_cont = ggml_is_contiguous(src1);
- if (src1_cont) {
+ if (!params->use_ref && src1_cont) {
for (int64_t i13 = 0; i13 < ne13; i13++)
for (int64_t i12 = 0; i12 < ne12; i12++)
if (!llamafile_sgemm(params,
@@ -1384,7 +1384,7 @@ UseGgmlGemm1:;
ggml_barrier(params->threadpool);
#if GGML_USE_LLAMAFILE
- if (src1->type != vec_dot_type) {
+ if (!params->use_ref && src1->type != vec_dot_type) {
const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
const size_t row_size = ggml_row_size(vec_dot_type, ne10);
diff --git a/ggml/src/ggml-cpu/llamafile/sgemm.cpp b/ggml/src/ggml-cpu/llamafile/sgemm.cpp
index 258bc9e2e..d28840a91 100644
--- a/ggml/src/ggml-cpu/llamafile/sgemm.cpp
+++ b/ggml/src/ggml-cpu/llamafile/sgemm.cpp
@@ -384,6 +384,80 @@ template <> inline __m256bh load(const float *p) {
}
#endif
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+template <typename T, typename U> T load_partial(const U *, int);
+template <typename T> T load_partial_u16(const void *, int);
+
+template <> inline __m128i load_partial_u16(const void *p, int n) {
+#if defined(__AVX512BW__) && defined(__AVX512VL__)
+ return _mm_maskz_loadu_epi16((1u << n) - 1, p);
+#else
+ const __m128i index = _mm_setr_epi32(0, 1, 2, 3);
+ const __m128i pairs = _mm_set1_epi32(n / 2);
+ __m128i v = _mm_castps_si128(_mm_maskload_ps((const float *)p, _mm_cmpgt_epi32(pairs, index)));
+ if (n & 1) {
+ uint16_t last;
+ memcpy(&last, (const char *)p + 2*(n - 1), sizeof(last));
+ v = _mm_or_si128(v, _mm_and_si128(_mm_cmpeq_epi32(pairs, index), _mm_set1_epi32(last)));
+ }
+ return v;
+#endif
+}
+
+template <> inline __m256 load_partial(const float *p, int n) {
+ const __m256 index = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7);
+ return _mm256_maskload_ps(p, _mm256_castps_si256(_mm256_cmp_ps(index, _mm256_set1_ps(n), _CMP_LT_OQ)));
+}
+
+#if defined(__F16C__)
+template <> inline __m256 load_partial(const ggml_fp16_t *p, int n) {
+ return _mm256_cvtph_ps(load_partial_u16<__m128i>(p, n));
+}
+#endif
+
+#if defined(__AVX2__) || defined(__AVX512F__)
+template <> inline __m256 load_partial(const ggml_bf16_t *p, int n) {
+ return _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(load_partial_u16<__m128i>(p, n)), 16));
+}
+#endif
+
+#if defined(__AVX512F__)
+template <> inline __m256i load_partial_u16(const void *p, int n) {
+#if defined(__AVX512BW__) && defined(__AVX512VL__)
+ return _mm256_maskz_loadu_epi16((1u << n) - 1, p);
+#else
+ const __m256i index = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+ const __m256i pairs = _mm256_set1_epi32(n / 2);
+ __m256i v = _mm256_maskload_epi32((const int *)p, _mm256_cmpgt_epi32(pairs, index));
+ if (n & 1) {
+ uint16_t last;
+ memcpy(&last, (const char *)p + 2*(n - 1), sizeof(last));
+ v = _mm256_or_si256(v, _mm256_and_si256(_mm256_cmpeq_epi32(pairs, index), _mm256_set1_epi32(last)));
+ }
+ return v;
+#endif
+}
+
+template <> inline __m512 load_partial(const float *p, int n) {
+ return _mm512_maskz_loadu_ps((1u << n) - 1, p);
+}
+
+template <> inline __m512 load_partial(const ggml_fp16_t *p, int n) {
+ return _mm512_cvtph_ps(load_partial_u16<__m256i>(p, n));
+}
+
+template <> inline __m512 load_partial(const ggml_bf16_t *p, int n) {
+ return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(load_partial_u16<__m256i>(p, n)), 16));
+}
+#endif
+
+#if defined(__AVX512BF16__)
+template <> inline __m512bh load_partial(const ggml_bf16_t *p, int n) {
+ return (__m512bh) _mm512_maskz_loadu_epi16((uint64_t(1) << n) - 1, p);
+}
+#endif
+#endif
+
#if defined(__riscv_v_intrinsic)
template <> inline vfloat32m1_t load(const float *p) {
return __riscv_vle32_v_f32m1(p, __riscv_vsetvlmax_e32m1());
@@ -492,8 +566,10 @@ class tinyBLAS {
}
bool matmul(int64_t m, int64_t n) {
+#if !defined(__AVX__) && !defined(__AVX2__) && !defined(__AVX512F__)
if (k % KN != 0)
return false;
+#endif
// compute RM for only need tile with size RM&RM-1
#if VECTOR_REGISTERS == 32
if (m % 16 == 0 && (m/16 >= params->nth)) {
@@ -548,7 +624,7 @@ class tinyBLAS {
template <int RM, int RN>
inline void gemm_bloc(int64_t ii, int64_t jj) {
D Cv[RN][RM] = {};
- for (int64_t l = 0; l < k; l += KN) {
+ for (int64_t l = 0; l + KN <= k; l += KN) {
// help compiler for op order.
if constexpr (RM <= RN) {
V Av[RM];
@@ -574,6 +650,21 @@ class tinyBLAS {
}
}
}
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+ const int64_t rem = k % KN;
+ if (rem != 0) {
+ V Av[RM];
+ for (int64_t i = 0; i < RM; ++i) {
+ Av[i] = load_partial<V>(A + lda * (ii + i) + k - rem, rem);
+ }
+ for (int64_t j = 0; j < RN; ++j) {
+ V Bv = load_partial<V>(B + ldb * (jj + j) + k - rem, rem);
+ for (int64_t i = 0; i < RM; ++i) {
+ Cv[j][i] = madd(Av[i], Bv, Cv[j][i]);
+ }
+ }
+ }
+#endif
for (int64_t j = 0; j < RN; ++j)
for (int64_t i = 0; i < RM; ++i)
C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 507fda75a..070d378f4 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10448,7 +10448,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6, 4096, 5120, {1, 1}, {1, 1}));
- // K not a multiple of 32
+ // SIMD K boundaries, including tails shorter than one vector.
+ for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) {
+ for (int k : {1, 7, 8, 9, 15, 16, 17, 31, 32, 33, 48}) {
+ test_cases.emplace_back(new test_mul_mat(type, type, 64, 17, k, {1, 1}, {1, 1}));
+ if (type != GGML_TYPE_F32) {
+ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, 17, k, {1, 1}, {1, 1}));
+ }
+ }
+ }
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 65, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 64, 32, 80, {1, 1}, {1, 1}));