CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-11-05 09:36:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

repack : optimize mul_mat_id path

Browse Source 
ggml-ci
This commit is contained in:
Georgi Gerganov
2025-07-28 15:19:04 +03:00
parent e2661edd24
commit 477d43988a
1 changed files with 154 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

229
ggml/src/ggml-cpu/repack.cpp
View File

@@ -1312,15 +1312,39 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
} }
case GGML_OP_MUL_MAT_ID: case GGML_OP_MUL_MAT_ID:
{ {
size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1])); const ggml_tensor * src0 = op->src[0];
size = GGML_PAD(size, sizeof(int64_t)); // + padding for next bloc. const ggml_tensor * src1 = op->src[1];
const ggml_tensor * dst = op;
const int64_t ne02 = op->src[0]->ne[2]; // n_as, n_expert GGML_TENSOR_BINARY_OP_LOCALS
const int64_t ne12 = op->src[1]->ne[2]; // n_tokens
const size_t sizeof_mmid_row_mapping = sizeof(int64_t); // src0 [n_embd, n_rows, n_expert]
// src1 [n_embd, n_expert_used, n_tokens]
// dst [n_rows, n_expert_used, n_tokens]
size += sizeof_mmid_row_mapping*ne02*(ne12 + 1); // htmp [n_embd, n_tokens, n_expert] F32
size_t size_htmp = ggml_row_size(GGML_TYPE_F32, ne00*ne12*ne02);
// hsrc1 [n_embd, n_tokens, n_expert]
size_t size_hsrc1 = ggml_row_size(PARAM_TYPE, ne00*ne12*ne02);
// hdst [n_rows, n_tokens, n_expert]
size_t size_hdst = ggml_row_size(GGML_TYPE_F32, ne01*ne12*ne02);
// htpe [n_expert]
size_t size_htpe = ggml_row_size(GGML_TYPE_I32, ne02);
// hids [n_expert*n_tokens]
size_t size_hids = ggml_row_size(GGML_TYPE_I32, ne02*ne12);
// + padding
size_htmp = GGML_PAD(size_htmp, sizeof(int64_t));
size_hsrc1 = GGML_PAD(size_hsrc1, sizeof(int64_t));
size_hdst = GGML_PAD(size_hdst, sizeof(int64_t));
size_htpe = GGML_PAD(size_htpe, sizeof(int64_t));
size_hids = GGML_PAD(size_hids, sizeof(int64_t));
size = size_htmp + size_hsrc1 + size_hdst + size_htpe + size_hids;
return true; return true;
} }
@@ -1446,76 +1470,112 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
GGML_ASSERT(src1->type == GGML_TYPE_F32); GGML_ASSERT(src1->type == GGML_TYPE_F32);
// row groups const int64_t ne20 = ids->ne[0];
const int n_ids = ids->ne[0]; // n_expert_used
const int n_as = ne02; // n_expert
const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10); // src0 [n_embd, n_rows, n_expert]
const size_t nbw2 = nbw1*ne11; // src1 [n_embd, n_expert_used', n_tokens]
const size_t nbw3 = nbw2*ne12; // src2 [n_expert_used, n_tokens]
// dst [n_rows, n_expert_used, n_tokens]
struct mmid_row_mapping { // htmp [n_embd, n_tokens, n_expert] F32
int32_t i1; size_t size_htmp = ggml_row_size(GGML_TYPE_F32, ne00*ne12*ne02);
int32_t i2;
};
GGML_ASSERT(params->wsize >= // hsrc1 [n_embd, n_tokens, n_expert]
(GGML_PAD(nbw3, sizeof(int64_t)) + size_t size_hsrc1 = ggml_row_size(PARAM_TYPE, ne00*ne12*ne02);
n_as*(ne12 + 1)*sizeof(mmid_row_mapping))
);
auto * wdata = (char *)params->wdata; // hdst [n_rows, n_tokens, n_expert]
auto * wdata_src1_end = (char *)wdata + GGML_PAD(nbw3, sizeof(int64_t)); size_t size_hdst = ggml_row_size(GGML_TYPE_F32, ne01*ne12*ne02);
// total of [n_as][ne12 + 1] elemets of type mmid_row_mapping (2*int32_t = int64_t) // htpe [n_expert]
auto * matrix_row_counts = (int64_t *) (wdata_src1_end); // [n_as] size_t size_htpe = ggml_row_size(GGML_TYPE_I32, ne02);
struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) (matrix_row_counts + n_as); // [n_as][ne12]
// src1: float32 => param type // hids [n_expert*n_tokens]
for (int64_t i12 = 0; i12 < ne12; ++i12) { size_t size_hids = ggml_row_size(GGML_TYPE_I32, ne02*ne12);
for (int64_t i11 = ith; i11 < ne11; i11 += nth) {
from_float((float *)((char *) src1->data + i12 * nb12 + i11 * nb11), // + padding
(void *) (wdata + i12 * nbw2 + i11 * nbw1), size_htmp = GGML_PAD(size_htmp, sizeof(int64_t));
ne10); size_hsrc1 = GGML_PAD(size_hsrc1, sizeof(int64_t));
size_hdst = GGML_PAD(size_hdst, sizeof(int64_t));
size_htpe = GGML_PAD(size_htpe, sizeof(int64_t));
size_hids = GGML_PAD(size_hids, sizeof(int64_t));
char * wdata_htmp = (char *) params->wdata;
char * wdata_hsrc1 = (char *) params->wdata + size_htmp;
char * wdata_hdst = (char *) params->wdata + size_htmp + size_hsrc1;
char * wdata_htpe = (char *) params->wdata + size_htmp + size_hsrc1 + size_hdst;
char * wdata_hids = (char *) params->wdata + size_htmp + size_hsrc1 + size_hdst + size_htpe;
const size_t nbht1 = ggml_row_size(GGML_TYPE_F32, ne00);
const size_t nbht2 = nbht1*ne12;
const size_t nbh11 = ggml_row_size(PARAM_TYPE, ne00);
const size_t nbh12 = nbh11*ne12;
const size_t nbh1 = ggml_row_size(GGML_TYPE_F32, ne01);
const size_t nbh2 = nbh1*ne12;
char * htmp = (char *)(wdata_htmp);
char * hsrc1 = (char *)(wdata_hsrc1);
char * hdst = (char *)(wdata_hdst);
int32_t * htpe = (int32_t *)(wdata_htpe);
int32_t * hids = (int32_t *)(wdata_hids);
for (int64_t i02 = ith; i02 < ne02; i02 += nth) {
htpe[i02] = 0;
}
// src1 (float32) => htmp (float32)
for (int64_t i12 = 0; i12 < ne12; ++i12) { // n_tokens
for (int64_t i20 = 0; i20 < ne20; ++i20) { // n_expert_used
// the selected expert
const int32_t i02 = *(const int32_t *) ((const char *) ids->data + i12*ids->nb[1] + i20*ids->nb[0]);
if (i02 % nth != ith) {
continue;
}
memcpy( htmp + i02*nbht2 + htpe[i02]*nbht1,
(char *) src1->data + i12*nb12 + (i20%ne11)*nb11,
ggml_row_size(GGML_TYPE_F32, ne10));
hids[i12*ne20 + i20] = i02*ne12 + htpe[i02];
htpe[i02]++;
} }
} }
#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ne12 + (i1)] // htmp (float32) => hsrc1 (param type)
for (int64_t i02 = 0; i02 < ne02; ++i02) { // n_expert
if (ith == 0) { if (i02 % nth != ith) {
// initialize matrix_row_counts continue;
memset(matrix_row_counts, 0, n_as * sizeof(int64_t));
// group rows by src0 matrix
for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
for (int32_t id = 0; id < n_ids; ++id) {
const int32_t i02 =
*(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
GGML_ASSERT(i02 >= 0 && i02 < n_as);
MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
matrix_row_counts[i02] += 1;
} }
const int64_t neh11 = htpe[i02];
for (int64_t i11 = 0; i11 < neh11 - neh11 % 4; i11 += 4) {
ggml_quantize_mat_t<INTER_SIZE, PARAM_TYPE>(
(float *) (htmp + i11*nbht1 + i02*nbht2),
(void *) (hsrc1 + i11*nbh11 + i02*nbh12), 4, ne10);
}
for (int64_t i11 = neh11 - neh11 % 4; i11 < neh11; i11 += 1) {
from_float(
(float *) (htmp + i11*nbht1 + i02*nbht2),
(void *) (hsrc1 + i11*nbh11 + i02*nbh12), ne10);
} }
} }
ggml_barrier(params->threadpool); ggml_barrier(params->threadpool);
// compute each matrix multiplication in sequence for (int64_t i02 = 0; i02 < ne02; ++i02) { // n_expert
for (int cur_a = 0; cur_a < n_as; ++cur_a) { const int64_t neh11 = htpe[i02];
const int64_t cne1 = matrix_row_counts[cur_a];
if (cne1 == 0) { if (neh11 == 0) {
continue; continue;
} }
const auto * src0_cur = (const char *) src0->data + cur_a*nb02; const auto * src0_cur = (const char *) src0->data + i02*nb02;
//const int64_t nr0 = ne01; // src0 rows int64_t src0_cur_start = ((ith )*ne01)/nth;
const int64_t nr1 = cne1; // src1 rows
int64_t src0_cur_start = (ith * ne01) / nth;
int64_t src0_cur_end = ((ith + 1)*ne01)/nth; int64_t src0_cur_end = ((ith + 1)*ne01)/nth;
src0_cur_start = (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start; src0_cur_start = (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start;
@@ -1525,26 +1585,49 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
return; return;
} }
for (int ir1 = 0; ir1 < nr1; ir1++) { #if 1
struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1); if (neh11 > 3) {
gemm<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
const int id = row_mapping.i1; // selected expert index (float *)(hdst + 0*nbh1 + i02*nbh2) + src0_cur_start, ne01,
const int64_t i11 = id % ne11;
const int64_t i12 = row_mapping.i2; // row index in src1
const int64_t i1 = id; // selected expert index
const int64_t i2 = i12; // row
const auto * src1_col = (const char *) wdata + (i11 * nbw1 + i12 * nbw2);
gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
(float *)((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, ne01,
src0_cur + src0_cur_start*nb01, src0_cur + src0_cur_start*nb01,
src1_col, 1, src0_cur_end - src0_cur_start); hsrc1 + 0*nbh11 + i02*nbh12, neh11 - neh11 % 4, src0_cur_end - src0_cur_start);
}
for (int64_t i11 = neh11 - neh11 % 4; i11 < neh11; ++i11) {
gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
(float *)(hdst + i11*nbh1 + i02*nbh2) + src0_cur_start, ne01,
src0_cur + src0_cur_start*nb01,
hsrc1 + i11*nbh11 + i02*nbh12, 1, src0_cur_end - src0_cur_start);
}
#else
for (int64_t i11 = 0; i11 < neh11; ++i11) {
gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
(float *)(hdst + i11*nbh1 + i02*nbh2) + src0_cur_start, ne01,
src0_cur + src0_cur_start * nb01,
hsrc1 + i11*nbh11 + i02*nbh12, 1, src0_cur_end - src0_cur_start);
}
#endif
}
ggml_barrier(params->threadpool);
for (int64_t i21 = 0; i21 < ne12; ++i21) { // n_tokens
for (int64_t i20 = 0; i20 < ne20; ++i20) { // n_expert_used
const int32_t idx = i21*ne20 + i20;
if (idx % nth != ith) {
continue;
}
const int32_t id = hids[idx];
const int ide = id/ne12;
const int idt = id%ne12;
memcpy(
(char *) dst->data + i20*nb1 + i21*nb2,
hdst + idt*nbh1 + ide*nbh2, ggml_row_size(GGML_TYPE_F32, ne01));
} }
} }
#undef MMID_MATRIX_ROW
} }
int repack(struct ggml_tensor * t, const void * data, size_t data_size) override { int repack(struct ggml_tensor * t, const void * data, size_t data_size) override {