happyz
/
llama.cpp
mirror of https://github.com/ggerganov/llama.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

ggml-blas: initial mmid impl 

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
This commit is contained in:
Aaron Teo 2025-12-11 20:51:02 +08:00
parent e4ae383317
commit f682374613
No known key found for this signature in database
1 changed files with 154 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

154
ggml/src/ggml-blas/ggml-blas.cpp
View File

@ -149,6 +149,143 @@ static void ggml_backend_blas_mul_mat(ggml_backend_blas_context * ctx, struct gg
}
}
static void ggml_backend_blas_mul_mat_id(ggml_backend_blas_context * ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0]; // weights
const ggml_tensor * src1 = dst->src[1]; // inputs
const ggml_tensor * src2 = dst->src[2]; // ids
GGML_TENSOR_TERNARY_OP_LOCALS
const ggml_type type = src0->type;
GGML_ASSERT(ne10 == ne00);
GGML_ASSERT(ne21 == ne12);
GGML_ASSERT(ne22 == 1 || ne22 == ne13);
GGML_ASSERT(src2->type == GGML_TYPE_I32);
GGML_ASSERT(nb00 == ggml_type_size(type));
GGML_ASSERT(nb10 == ggml_type_size(src1->type));
GGML_ASSERT(dst->type == GGML_TYPE_F32);
GGML_ASSERT(nb0 == sizeof(float));
GGML_ASSERT(nb0 <= nb1 && nb1 <= nb2 && nb2 <= nb3);
const int64_t n_used = (int64_t)ne20;
GGML_ASSERT(n_used <= ne02);
const int64_t ne_plane = ne01 * ne00;
const size_t desired_wsize = (type == GGML_TYPE_F32) ? 0 : ne03 * ne02 * ne_plane * sizeof(float);
if (ctx->work_size < desired_wsize) {
ctx->work_data.reset(new char[desired_wsize]);
ctx->work_size = desired_wsize;
}
void * wdata = ctx->work_data.get();
if (type != GGML_TYPE_F32) {
const auto * type_traits = ggml_get_type_traits(type);
ggml_to_float_t to_float = type_traits->to_float;
for (int64_t i03 = 0; i03 < ne03; ++i03) {
for (int64_t i02 = 0; i02 < ne02; ++i02) {
const void * x = (char *)src0->data + i02*nb02 + i03*nb03;
float * wplane = (float *)wdata + i02*ne_plane + i03*ne02*ne_plane;
const int min_cols_per_thread = 4096;
const int min_rows_per_thread = std::max((int)(min_cols_per_thread / ne00), 1);
const int n_threads = std::max(std::min(ctx->n_threads, (int)(ne01 / min_rows_per_thread)), 1);
#ifdef GGML_USE_OPENMP
#pragma omp parallel for num_threads(n_threads)
for (int64_t i01 = 0; i01 < ne01; ++i01) {
to_float((const char *)x + i01*nb01, wplane + i01*ne00, ne00);
}
#else
for (int i = 1; i < n_threads; i++) {
const int64_t start = i * ne01/n_threads;
const int64_t end = (i + 1) * ne01/n_threads;
if (start < end) {
ctx->tasks.push_back(std::async(std::launch::async, [=]() {
for (int64_t i01 = start; i01 < end; ++i01) {
to_float((const char *)x + i01*nb01, wplane + i01*ne00, ne00);
}
}));
}
}
{
const int64_t start = 0;
const int64_t end = ne01/n_threads;
for (int64_t i01 = start; i01 < end; i01++) {
to_float((const char *)x + i01*nb01, wplane + i01*ne00, ne00);
}
}
#endif
}
}
#ifndef GGML_USE_OPENMP
for (auto & task: ctx->tasks) {
task.get();
}
ctx->tasks.clear();
#endif
}
#ifdef OPENBLAS_VERSION
openblas_set_num_threads(ctx->n_threads);
#endif
#ifdef GGML_BLAS_USE_BLIS
bli_thread_set_num_threads(ctx->n_threads);
#endif
#ifdef GGML_BLAS_USE_NVPL
nvpl_blas_set_num_threads(ctx->n_threads);
#endif
for (int64_t i13 = 0; i13 < ne13; ++i13) {
for (int64_t j = 0; j < ne12; ++j) {
const int64_t ids_batch_index = (ne22 > 1 ? i13 : 0);
const int32_t * ids_row = (const int32_t *)((char *)src2->data + ids_batch_index*nb22 + j*nb21);
float * out_ptr = (float *)((char *)dst->data + i13*nb3 + j*nb2);
for (int iE = 0; iE < n_used; ++iE) {
const int expert_id = ids_row[iE];
GGML_ASSERT(expert_id < ne02);
const float * wmat;
if (type == GGML_TYPE_F32) {
wmat = (const float *)((char *)src0->data + expert_id*nb02);
} else {
wmat = (const float *)((char *)wdata + expert_id * ne_plane * sizeof(float));
}
if (ne03 > 1) {
int64_t w_batch_index = (ne03 == ne13 ? i13 : 0);
wmat = (const float *)((char *)wdata + (w_batch_index * ne02 + expert_id) * ne_plane * sizeof(float));
}
const float * inp = (const float *)((char *)src1->data
+ ((ne11 == 1 ? 0 : iE) * nb11)
+ j * nb12 + i13 * nb13);
if (iE == 0) {
cblas_sgemv(CblasRowMajor, CblasNoTrans, (int)ne01, (int)ne00,
1.0f, wmat, (int)ne00,
inp, 1,
0.0f,
out_ptr, 1);
} else {
cblas_sgemv(CblasRowMajor, CblasNoTrans, (int)ne01, (int)ne00,
1.0f, wmat, (int)ne00,
inp, 1,
1.0f,
out_ptr, 1);
}
}
}
}
}
static void ggml_backend_blas_out_prod(ggml_backend_blas_context * ctx, struct ggml_tensor * dst) {
const struct ggml_tensor * src0 = dst->src[0];
const struct ggml_tensor * src1 = dst->src[1];
@ -235,6 +372,10 @@ static enum ggml_status ggml_backend_blas_graph_compute(ggml_backend_t backend,
ggml_backend_blas_mul_mat(ctx, node);
break;
case GGML_OP_MUL_MAT_ID:
ggml_backend_blas_mul_mat_id(ctx, node);
break;
case GGML_OP_OUT_PROD:
ggml_backend_blas_out_prod(ctx, node);
break;
@ -418,6 +559,19 @@ static bool ggml_backend_blas_device_supports_op(ggml_backend_dev_t dev, const s
(src0->type == GGML_TYPE_F32 || ggml_get_type_traits(src0->type)->to_float != NULL);
}
case GGML_OP_MUL_MAT_ID:
{
const struct ggml_tensor * src0 = op->src[0];
const struct ggml_tensor * src1 = op->src[1];
const struct ggml_tensor * src2 = op->src[2];
// GGML_LOG_INFO("%s: op=GGML_OP_MUL_MAT_ID src0_type=%s src1_type=%s src2_type=%s ne0=%lld ne1=%lld ne2=%lld ne3=%lld\n",
// __func__, ggml_type_name(src0->type), ggml_type_name(src1->type), ggml_type_name(src2->type),
// op->ne[0], op->ne[1], op->ne[2], op->ne[3]);
return src2->type == GGML_TYPE_I32;
}
case GGML_OP_OUT_PROD:
return op->src[0]->type == GGML_TYPE_F32 &&
op->src[1]->type == GGML_TYPE_F32 &&