Skip to content

ggml-cpu: improve --n-cpu-moe TG performance #20596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
am17an wants to merge 4 commits into
base: master
Choose a base branch
from
Open

ggml-cpu: improve --n-cpu-moe TG performance #20596

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ae0023a
ggml-cpu: improve `--n-cpu-moe` TG performance
am17an Mar 15, 2026
873ef8f
rename, add cache line padding
am17an Mar 16, 2026
ddb029c
handle broadcast
am17an Mar 16, 2026
5d03c12
use work-stealing
am17an Mar 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 152 additions & 2 deletions ggml/src/ggml-cpu/ggml-cpu.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1681,6 +1681,129 @@ static void ggml_compute_forward_mul_mat_id(
}
}

static void ggml_compute_forward_fused_moe_silu(
const struct ggml_compute_params * params,
struct ggml_tensor * node0,
struct ggml_tensor * node1,
struct ggml_tensor * glu_node) {

const struct ggml_tensor * weights_gate;
const struct ggml_tensor * weights_up;

if (glu_node->src[0] == node0) {
weights_gate = node0->src[0];
weights_up = node1->src[0];
} else {
weights_gate = node1->src[0];
weights_up = node0->src[0];
}

const struct ggml_tensor * src1 = node0->src[1];
const struct ggml_tensor * ids = node0->src[2];

const int64_t ne00 = weights_gate->ne[0];
const int64_t ne01 = weights_gate->ne[1];

const size_t gate_nb01 = weights_gate->nb[1];
const size_t gate_nb02 = weights_gate->nb[2];
const size_t up_nb01 = weights_up->nb[1];
const size_t up_nb02 = weights_up->nb[2];

const int64_t ne10 = src1->ne[0];
const int64_t ne11 = src1->ne[1];
const size_t nb11 = src1->nb[1];

const size_t glu_nb1 = glu_node->nb[1];

const int ith = params->ith;
const int nth = params->nth;

const enum ggml_type type = weights_gate->type;

ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot;
enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float;

const int n_ids = ids->ne[0]; // n_expert_used
const int n_as = weights_gate->ne[2]; // n_experts

const size_t row_size = ggml_row_size(vec_dot_type, ne10);

void * wdata_cur = params->wdata;

const char * src1_q = (const char *) src1->data;
if (src1->type != vec_dot_type) {
char * quant_base = (char *) incr_ptr_aligned(&wdata_cur,
(ggml_row_size(vec_dot_type, ggml_nelements(src1)) + CACHE_LINE_SIZE) * nth, sizeof(int64_t));

GGML_ASSERT(src1->type == GGML_TYPE_F32);
char * wdata = quant_base + ith * (ne11 * row_size + CACHE_LINE_SIZE);
for (int64_t i11 = 0; i11 < ne11; ++i11) {
from_float((float *)((char *) src1->data + i11*nb11),
(void *)(wdata + i11*row_size),
ne10);
}
src1_q = wdata;
}

incr_ptr_aligned(&wdata_cur, n_as * sizeof(int64_t), sizeof(int64_t));
incr_ptr_aligned(&wdata_cur, n_as * ids->ne[0] * ids->ne[1] * sizeof(struct mmid_row_mapping), sizeof(int64_t));

char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as]
incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE);

if (ith == 0) {
for (int id = 0; id < n_ids; ++id) {
const int32_t expert_idx = *(const int32_t *) ((const char *) ids->data + id*ids->nb[0]);
atomic_int * ctr = (atomic_int *)(atomic_current_chunk + expert_idx);
atomic_store_explicit(ctr, nth, memory_order_relaxed);
}
}

ggml_barrier(params->threadpool);

const int64_t nr0 = ne01;
const int chunk_size = 64;
const bool disable_chunking = ggml_is_numa();

int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size;
if (nchunk0 < (int64_t)(nth * 4) || disable_chunking) {
nchunk0 = nth;
}
const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0;

for (int id = 0; id < n_ids; ++id) {
const int32_t expert_idx = *(const int32_t *) ((const char *) ids->data + id*ids->nb[0]);

const char * gate_cur = (const char *) weights_gate->data + expert_idx * gate_nb02;
const char * up_cur = (const char *) weights_up->data + expert_idx * up_nb02;
const char * src1_col = src1_q;

float * glu_col = (float *) ((char *) glu_node->data + id*glu_nb1);

atomic_int * current_chunk_ctr = (atomic_int *)(atomic_current_chunk + expert_idx);

int current_chunk = ith;
while (current_chunk < nchunk0) {
const int64_t ir0_start = dr0 * current_chunk;
const int64_t ir0_end = MIN(ir0_start + dr0, nr0);

for (int64_t ir0 = ir0_start; ir0 < ir0_end; ++ir0) {
float gate_val, up_val;
vec_dot(ne00, &gate_val, 0, gate_cur + ir0*gate_nb01, 0, src1_col, 0, 1);
vec_dot(ne00, &up_val, 0, up_cur + ir0*up_nb01, 0, src1_col, 0, 1);
glu_col[ir0] = ggml_silu_f32(gate_val) * up_val;
}

if (nth >= nchunk0) {
break;
}

current_chunk = atomic_fetch_add_explicit(current_chunk_ctr, 1, memory_order_relaxed);
}
}
}

/////////////////////////////////

static void ggml_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor) {
Expand Down Expand Up @@ -2809,7 +2932,12 @@ struct ggml_cplan ggml_graph_plan(
const int n_as = src0->ne[2];
// src1
if (src1->type != vec_dot_type) {
cur += ggml_row_size(vec_dot_type, ggml_nelements(src1)) + sizeof(int64_t);
size_t quant_buf = ggml_row_size(vec_dot_type, ggml_nelements(src1));
// fused MoE path: each thread needs its own quantization buffer + cache line padding
if (src1->ne[2] == 1) {
quant_buf = (quant_buf) * n_tasks;
}
cur += quant_buf + sizeof(int64_t);
}
// matrix_row_counts
cur += n_as * sizeof(int64_t) + sizeof(int64_t);
Expand Down Expand Up @@ -2985,7 +3113,29 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
continue;
}

ggml_compute_forward(&params, node);
// Try fusion: MUL_MAT_ID + MUL_MAT_ID + GLU
int fused_nodes = 0;
if (node->op == GGML_OP_MUL_MAT_ID) {
enum ggml_op fuse_ops[3] = {GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU};
int outputs[1] = {node_n + 2};
if (ggml_can_fuse_subgraph(cgraph, node_n, 3, fuse_ops, outputs, 1)) {
struct ggml_tensor * node1 = cgraph->nodes[node_n + 1];
struct ggml_tensor * glu = cgraph->nodes[node_n + 2];
// Fused path for `--n-cpu-moe` when n_tokens = 1

@ggerganov ggerganov Mar 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we do it only for n_tokens == 1?

@am17an am17an Mar 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Single token is a special case because you don't need to create the row mapping. Also in this I let each thread quantize the activation to remove a barrier. The barrier cost was around 30% of the total run time in the cpu part, with this change it's about 7%

@ggerganov ggerganov Mar 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, that's unexpected for the barrier to be so expensive. I'll need to double-check - do you have a patch that I can apply to test the barrier path?

Btw, you should space the wdata with CACHE_LINE_SIZE_F32 to avoid false sharing.

@max-krasnyansky max-krasnyansky Mar 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have test-barrier that measures the overhead with tiny graphs.
@am17an can you share what numbers you get from that.
Ideally with and without OMP (GGML_OPENMP=OFF)

@am17an am17an Mar 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ggerganov on master if you do

CUDA_VISIBLE_DEVICES=4 perf record ./build/bin/llama-bench -m /opt/models/Qwen3.5-35B-A3B-Q4_K_S.gguf -fa 1 -n 32 -r 20 -p 0 -ncmoe 99 -t 8

Image

@max-krasnyansky

with OpenMP (the default build)

graph-compute with
n_threads: 4
  n_nodes: 2000
 n_rounds: 100
graph-compute took 516226 usec
5162.26 usec per-iter
2581.13 nsec per-node
graph-compute with
n_threads: 4
  n_nodes: 4
 n_rounds: 10000
graph-compute with
gf0 n_nodes: 4
gf1 n_nodes: 8
  n_threads: 4
   n_rounds: 1000

With OpenMP=OFF

graph-compute with
 n_threads: 4
   n_nodes: 2000
  n_rounds: 100
graph-compute took 479589 usec
 4795.89 usec per-iter
 2397.95 nsec per-node
graph-compute with
 n_threads: 4
   n_nodes: 4
  n_rounds: 10000
graph-compute with
 gf0 n_nodes: 4
 gf1 n_nodes: 8
   n_threads: 4
    n_rounds: 1000

@am17an am17an Mar 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also since this barrier cost is non-trivial, it has a meaningful effect for this case, when there are too many threads contending. Maybe we should limit the n_tasks based on some heuristic. Can someone else also confirm these results (I'm not sure if NUMA is playing a role here). For reference all results on a AMD EPYC 7742 64-Core Processor

CUDA_VISIBLE_DEVICES=4 ./build/bin/llama-bench -m /opt/models/qwen3_next_80b_a3b_instruct-iq4_xs.gguf -fa 1 -n 32 -r 10 -p 0  -ncmoe 999 -t 8,16,32,48,64
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 32111 MiB):
  Device 0: NVIDIA GeForce RTX 5090, compute capability 12.0, VMM: yes, VRAM: 32111 MiB (31604 MiB free)
| model                          |       size |     params | backend    | ngl | threads | fa |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | -: | --------------: | -------------------: |
| qwen3next 80B.A3B IQ4_XS - 4.25 bpw |  39.67 GiB |    79.67 B | CUDA       |  99 |       8 |  1 |            tg32 |         30.85 ± 0.48 |
| qwen3next 80B.A3B IQ4_XS - 4.25 bpw |  39.67 GiB |    79.67 B | CUDA       |  99 |      16 |  1 |            tg32 |         29.81 ± 0.45 |
| qwen3next 80B.A3B IQ4_XS - 4.25 bpw |  39.67 GiB |    79.67 B | CUDA       |  99 |      32 |  1 |            tg32 |         28.76 ± 0.58 |
| qwen3next 80B.A3B IQ4_XS - 4.25 bpw |  39.67 GiB |    79.67 B | CUDA       |  99 |      48 |  1 |            tg32 |         27.48 ± 0.30 |
| qwen3next 80B.A3B IQ4_XS - 4.25 bpw |  39.67 GiB |    79.67 B | CUDA       |  99 |      64 |  1 |            tg32 |         26.19 ± 0.35 |

@ggerganov ggerganov Mar 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll take a look later today. Btw, I was rather thinking to compare the fused version:

  • no barrier + all threads reading src1 (current PR)
  • barrier + 1 thread reading src1

@am17an am17an Mar 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the patch here 0c0cf6f

@max-krasnyansky max-krasnyansky Mar 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is what I'm getting with OpenMP=OFF

AMD EPYC 7543
  5177.86 usec per-iter
  2588.93 nsec per-node

AMD Ryzen 395+
   1559.2 usec per-iter
   779.601 nsec per-node

Mac M4 Pro
  1338.91 usec per-iter
  669.456 nsec per-node 
 
Galaxy S24+
  1873.33 usec per-iter
  936.669 nsec per-node

Your numbers are a bit high but inline with the AMD EPYC.

With OpenMP=OFF

 4795.89 usec per-iter
 2397.95 nsec per-node

Also from your profile results it looks like it's not really the barrier itself that is expensive.
Most of the cycles are spent in the barrier_wait.
That just means that some of the cores are completing their chunks faster than the others cores.
Maybe our chunking is suboptimal in this case?
We recently updated MUL_MAT and FA to improve dynamic chunking.
Perhaps, it's a similar issue here.

@am17an am17an Mar 17, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@max-krasnyansky that profile is from the master branch. MUL MAT ID also has the same chunked approach if I'm not wrong

if (node->src[1] == node1->src[1] && node->src[2] == node1->src[2] &&
ggml_nrows(node->src[1]) == 1 &&
ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU) {
ggml_compute_forward_fused_moe_silu(&params, node, node1, glu);
fused_nodes = 2;
}
}
}

if (fused_nodes == 0) {
ggml_compute_forward(&params, node);
}

node_n += fused_nodes;

if (state->ith == 0 && cplan->abort_callback &&
cplan->abort_callback(cplan->abort_callback_data)) {
Expand Down
Loading